From 0fc2f0e5ca3504384f3542bc2624d18dda5063d8 Mon Sep 17 00:00:00 2001 From: Michal Harakal <471011+michalharakal@users.noreply.github.com> Date: Thu, 5 Nov 2026 18:35:52 +0100 Subject: [PATCH] fix(compile-hlo): read block-quant weights through the packed view BlockQuantMatmulConverter read Q8_0 / Q4_0 blocks from heap packedData only. Weights from the GGUF loader are off-heap BufferPackedTensorData, whose packedData throws, so exporting a real checkpoint failed. The converter now decodes the canonical ggml blocks (Q8_0: f16 d + 32 i8; Q4_0: f16 d + 16 nibble bytes) by copying rows out of packedView.storage, which works for heap and loader-backed weights alike. Found exporting the full FunctionGemma decode step (127 folded projections). Co-Authored-By: Claude Opus 5.5 --- CHANGELOG.md | 1 + .../converters/BlockQuantMatmulConverter.kt | 53 +++++++++++++------ .../compile/hlo/BlockQuantMatmulExportTest.kt | 22 ++++++++ 3 files changed, 61 insertions(+), 15 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index f85744a7..d65c2000 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,7 @@ ### Fixed +- **Block-quant export reads loader-backed weights** (`skainet-compile-hlo`): `BlockQuantMatmulConverter` read Q8_0 / Q4_0 blocks only from heap `packedData`, so a weight from the GGUF loader (`BufferPackedTensorData`, whose `packedData` throws) failed the export. It now reads the ggml blocks through the weight's packed view, which works for both. - **StableHLO export maps `Int4` to `i4` and fails on a dtype without an element type** (`skainet-compile-hlo`): `TypeMapper.mapDType` turned every name it did not know, `Int4` and `Ternary` among them, into `f32`, so a typed `Int4` or `Ternary` tensor exported as a module whose types disagreed with its data. `Int4`/`I4` now map to `i4`; an unset dtype (`""`, `unknown`, what the DSL records when none is given) stays `f32`; any other unknown name throws `IllegalArgumentException` naming the dtype. Quantized weights are unaffected: they reach the emitter as dense tensors with a storage encoding. ### Added diff --git a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt index 19d6e0c7..f67295b9 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonMain/kotlin/sk/ainet/compile/hlo/converters/BlockQuantMatmulConverter.kt @@ -5,8 +5,6 @@ import sk.ainet.compile.hlo.ConversionResult import sk.ainet.compile.hlo.StableHloOperationConverter import sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass import sk.ainet.lang.graph.GraphNode -import sk.ainet.lang.tensor.data.Q4_0TensorData -import sk.ainet.lang.tensor.data.Q8_0TensorData import sk.ainet.lang.tensor.storage.PackedBlockStorage import sk.ainet.lang.tensor.storage.TensorEncoding @@ -203,37 +201,49 @@ internal class BlockQuantWeight( val n = storage.shape[0] val k = storage.shape[1] require(k % BLOCK == 0) { "block-quant weight in = $k is not a multiple of $BLOCK" } - val (bits, scaleOf, codeOf) = when (storage.encoding) { - TensorEncoding.Q8_0 -> { - val q = storage as? Q8_0TensorData ?: throw IllegalArgumentException("Q8_0 storage ${storage::class.simpleName} has no block accessors") - Triple(8, q::getBlockScale, { blk: Int, j: Int -> q.getCode(blk, j).toInt() }) - } - TensorEncoding.Q4_0 -> { - val q = storage as? Q4_0TensorData ?: throw IllegalArgumentException("Q4_0 storage ${storage::class.simpleName} has no block accessors") - Triple(4, q::getBlockScale, { blk: Int, j: Int -> q.getCode(blk, j).toInt() - 8 }) - } + // Canonical ggml blocks, read through the storage behind the packed view: heap and + // off-heap (mapped GGUF) weights alike. Q8_0: f16 d + 32 int8 codes. Q4_0: f16 d + 16 + // bytes, low nibbles elements 0..15, high nibbles 16..31, value (q - 8) * d. + val (bits, bytesPerBlock) = when (storage.encoding) { + TensorEncoding.Q8_0 -> 8 to 34 + TensorEncoding.Q4_0 -> 4 to 18 else -> throw IllegalArgumentException("block-quant matmul supports Q8_0 and Q4_0, not ${storage.encoding.name}") } + val view = storage.packedView + val source = view.storage + val base = view.layout.offsetBytes val nB = k / BLOCK val codes = ByteArray(nB * BLOCK * n) val scales = ShortArray(nB * n) var zeroPoints: ShortArray? = null + val row = ByteArray(nB * bytesPerBlock) + val block = IntArray(BLOCK) for (o in 0 until n) { + source.copyInto(row, 0, base + o.toLong() * row.size, row.size) for (b in 0 until nB) { - val blk = o * nB + b - val d = scaleOf(blk) + val at0 = b * bytesPerBlock + val d = halfToFloat((row[at0].toInt() and 0xFF) or ((row[at0 + 1].toInt() and 0xFF) shl 8)) + if (bits == 8) { + for (j in 0 until BLOCK) block[j] = row[at0 + 2 + j].toInt() + } else { + for (j in 0 until 16) { + val q = row[at0 + 2 + j].toInt() and 0xFF + block[j] = (q and 0x0F) - 8 + block[j + 16] = (q ushr 4) - 8 + } + } val at = b * n + o when { d > 0f -> { scales[at] = bf16(d) - for (j in 0 until BLOCK) codes[(b * BLOCK + j) * n + o] = codeOf(blk, j).toByte() + for (j in 0 until BLOCK) codes[(b * BLOCK + j) * n + o] = block[j].toByte() } d < 0f -> { // (c - 0) * d == (-1 - c - (-1)) * -d, and -1 - c stays in range for i8 and i4. scales[at] = bf16(-d) val zp = zeroPoints ?: ShortArray(nB * n).also { zeroPoints = it } zp[at] = BF16_MINUS_ONE - for (j in 0 until BLOCK) codes[(b * BLOCK + j) * n + o] = (-1 - codeOf(blk, j)).toByte() + for (j in 0 until BLOCK) codes[(b * BLOCK + j) * n + o] = (-1 - block[j]).toByte() } // An all-zero block: zero codes, and a scale the kernel accepts (it rejects 0). else -> scales[at] = BF16_ONE @@ -243,6 +253,19 @@ internal class BlockQuantWeight( return BlockQuantWeight(k, n, bits, codes, scales, zeroPoints) } + /** IEEE half bits → FP32. */ + fun halfToFloat(h: Int): Float { + val sign = (h ushr 15) and 1 + val exp = (h ushr 10) and 0x1F + val mant = h and 0x3FF + val magnitude = when (exp) { + 0 -> mant * 5.9604645e-8f // subnormal: mant * 2^-24 + 0x1F -> if (mant == 0) Float.POSITIVE_INFINITY else Float.NaN + else -> Float.fromBits(((exp + 112) shl 23) or (mant shl 13)) + } + return if (sign == 1) -magnitude else magnitude + } + /** FP32 → bf16 bits, rounded to nearest even. */ fun bf16(value: Float): Short { val bits = value.toRawBits() diff --git a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BlockQuantMatmulExportTest.kt b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BlockQuantMatmulExportTest.kt index 69210576..273447f6 100644 --- a/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BlockQuantMatmulExportTest.kt +++ b/skainet-compile/skainet-compile-hlo/src/commonTest/kotlin/sk/ainet/compile/hlo/BlockQuantMatmulExportTest.kt @@ -5,6 +5,9 @@ import sk.ainet.compile.opt.passes.BlockQuantMatmulFusionPass import sk.ainet.lang.graph.ComputeGraph import sk.ainet.lang.graph.DefaultGraphExecutionContext import sk.ainet.lang.tape.toComputeGraph +import sk.ainet.lang.memory.PackedBlockDecoder +import sk.ainet.lang.memory.Storage +import sk.ainet.lang.memory.TensorView import sk.ainet.lang.tensor.Shape import sk.ainet.lang.tensor.Tensor import sk.ainet.lang.tensor.data.Q4_0Quantizer @@ -127,6 +130,25 @@ class BlockQuantMatmulExportTest { } } + @Test + fun offHeapStyleStorageReadsTheSameBlocks() { + // A loader-backed weight (e.g. a mapped GGUF) has no heap packedData, only the storage behind + // its packed view. + for (heap in listOf(Q8_0Quantizer.quantize(weights(7), Shape(n, k)), Q4_0Quantizer.quantize(weights(8), Shape(n, k)))) { + val bytes = heap.packedData.copyOf() + val viewOnly = object : PackedBlockStorage by heap { + override val packedData: ByteArray get() = throw UnsupportedOperationException("off-heap") + override val packedStorage: Storage get() = Storage.Heap.wrap(bytes, mutable = false) + override val packedView: TensorView + get() = TensorView.packed(packedStorage, shape, encoding, PackedBlockDecoder(heap)) + } + val a = BlockQuantWeight.of(heap) + val b = BlockQuantWeight.of(viewOnly) + assertTrue(a.codes.contentEquals(b.codes) && a.scales.contentEquals(b.scales), "${heap.encoding.name} differs") + assertEquals(a.zeroPoints?.toList(), b.zeroPoints?.toList()) + } + } + @Test fun allZeroBlockGetsAScaleTheKernelAccepts() { val w = BlockQuantWeight.of(Q8_0Quantizer.quantize(FloatArray(2 * 32), Shape(2, 32)))