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 6757ae6ee..b6594f7a7 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 @@ -289,6 +289,7 @@ public class sk/ainet/exec/tensor/ops/DefaultCpuOpsBase : sk/ainet/lang/tensor/o public fun relu (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; protected final fun requireSameDType (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;)V public fun reshape (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Shape;)Lsk/ainet/lang/tensor/Tensor; + protected final fun rewrapFeedOrderWeight (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun rsubScalar (Ljava/lang/Number;Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; public fun scaledDotProductAttention (Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;Lsk/ainet/lang/tensor/Tensor;FZ)Lsk/ainet/lang/tensor/Tensor; public fun sigmoid (Lsk/ainet/lang/tensor/Tensor;)Lsk/ainet/lang/tensor/Tensor; 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 d19d8e384..15e2a531a 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 @@ -1002,6 +1002,36 @@ public open class DefaultCpuOpsBase(protected val dataFactory: TensorDataFactory "no packed relayout for ${weight.data::class.simpleName}", ) + /** + * A weight already stored in kernel feed order, relabelled as the `[in, out]` tensor the packed + * kernels take — **sharing the same bytes** (#1120). + * + * This is the payoff of letting packed storage declare its order. `transposePackedWeight` is an + * O(bytes) permutation run once per weight and cached; when the loader already produced feed + * order there is nothing to permute, and all that is needed is the other shape label over the + * same array. `null` when [tensor] is not a feed-order packed weight. + */ + @Suppress("UNCHECKED_CAST") + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + protected fun rewrapFeedOrderWeight(tensor: Tensor): Tensor? { + if (tensor.shape.rank != 2) return null + val packed = tensor.data as? sk.ainet.lang.tensor.storage.PackedBlockStorage ?: return null + if (packed.blockOrder != sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR) return null + val swapped = Shape(tensor.shape[1], tensor.shape[0]) + val bytes = packed.packedData + val relabelled: TensorData = when (tensor.data) { + is Q4_KTensorData -> Q4_KBlockTensorData(swapped, bytes) as TensorData + is Q5_KTensorData -> Q5_KBlockTensorData(swapped, bytes) as TensorData + is Q6_KTensorData -> Q6_KBlockTensorData(swapped, bytes) as TensorData + is Q5_1TensorData -> Q5_1BlockTensorData(swapped, bytes) as TensorData + is Q5_0TensorData -> Q5_0BlockTensorData(swapped, bytes) as TensorData + is Q8_0TensorData -> Q8_0BlockTensorData(swapped, bytes) as TensorData + is Q4_0TensorData -> Q4_0BlockTensorData(swapped, bytes) as TensorData + else -> return null + } + return newTensor(relabelled, tensor.dtype, tensor) + } + @Suppress("UNCHECKED_CAST") private fun transposePackedWeight(tensor: Tensor): Tensor? { val rank = tensor.shape.rank diff --git a/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/FeedOrderWeightNoCopyTest.kt b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/FeedOrderWeightNoCopyTest.kt new file mode 100644 index 000000000..5fd8712a3 --- /dev/null +++ b/skainet-backends/skainet-backend-cpu/src/commonTest/kotlin/sk/ainet/exec/tensor/ops/FeedOrderWeightNoCopyTest.kt @@ -0,0 +1,100 @@ +package sk.ainet.exec.tensor.ops + +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.data.Q8_0BlockTensorData +import sk.ainet.lang.tensor.data.TensorData +import sk.ainet.lang.tensor.matmulWeightTransposed +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.types.FP32 +import kotlin.math.abs +import kotlin.test.Test +import kotlin.test.assertEquals +import kotlin.test.assertSame +import kotlin.test.assertTrue + +/** + * #1120: a weight the loader already put in kernel feed order costs nothing to use. + * + * The point of declaring block order is not tidiness — it is that the O(bytes) relayout #1096 runs + * once per weight stops running at all. Asserted structurally rather than by timing: the bytes the + * kernels are handed must be *the same array*, not an equal one. + */ +@OptIn(ExperimentalMemoryApi::class) +class FeedOrderWeightNoCopyTest { + + private val ctx = DirectCpuExecutionContext() + private val outDim = 32 + private val inDim = 96 // three blocks per row: the orders differ (#968) + + private fun bytes(): ByteArray { + val out = ByteArray(outDim * (inDim / 32) * 34) + var seed = 7 + for (b in 0 until outDim * (inDim / 32)) { + val base = b * 34 + out[base] = 0x00; out[base + 1] = 0x3C + for (i in 0 until 32) { + seed = seed * 1103515245 + 12345 + out[base + 2 + i] = ((seed ushr 16) % 9 - 4).toByte() + } + } + return out + } + + @Suppress("UNCHECKED_CAST") + private fun weight(order: BlockOrder, payload: ByteArray): Tensor = + ctx.fromData( + Q8_0BlockTensorData(Shape(outDim, inDim), payload, order) as TensorData, + FP32::class, + ) + + @Test + fun `a feed-order weight reaches the kernels without its bytes being copied`() { + val payload = bytes() + val w = weight(BlockOrder.INPUT_BLOCK_MAJOR, payload) + + val x = ctx.fromFloatArray(Shape(1, inDim), FP32::class, FloatArray(inDim) { (it % 13) * 0.0625f }) + val result = x.matmulWeightTransposed(w) + assertEquals(Shape(1, outDim), result.shape) + + // The claim: the same array, still. A relayout would have produced a different one. + assertSame(payload, (w.data as PackedBlockStorage).packedData, "the weight's own bytes must be untouched") + } + + @Test + fun `a feed-order weight and the canonical weight it came from give the same product`() { + val canonical = bytes() + val canonicalWeight = weight(BlockOrder.ROW_MAJOR, canonical) + + // Permute canonically-ordered blocks into feed order by hand: block (o, b) moves from + // o * blocksPerRow + b to b * rows + o. + val blocksPerRow = inDim / 32 + val feed = ByteArray(canonical.size) + for (o in 0 until outDim) { + for (b in 0 until blocksPerRow) { + canonical.copyInto(feed, (b * outDim + o) * 34, (o * blocksPerRow + b) * 34, (o * blocksPerRow + b + 1) * 34) + } + } + val feedWeight = weight(BlockOrder.INPUT_BLOCK_MAJOR, feed) + + assertTrue(!canonical.contentEquals(feed), "the two orders must differ here or the test is vacuous") + assertTrue( + (canonicalWeight.data as PackedBlockStorage).toFloatArray() + .contentEquals((feedWeight.data as PackedBlockStorage).toFloatArray()), + "different bytes, same matrix", + ) + + val x = ctx.fromFloatArray(Shape(1, inDim), FP32::class, FloatArray(inDim) { (it % 13) * 0.0625f }) + val fromCanonical = x.matmulWeightTransposed(canonicalWeight).data.copyToFloatArray() + val fromFeed = x.matmulWeightTransposed(feedWeight).data.copyToFloatArray() + for (o in fromCanonical.indices) { + assertTrue( + abs(fromCanonical[o] - fromFeed[o]) <= 1e-3f * maxOf(1.0f, abs(fromCanonical[o])), + "output[$o]: canonical ${fromCanonical[o]} vs feed ${fromFeed[o]}", + ) + } + } +} 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 0f5992404..e62dde190 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 @@ -192,6 +192,10 @@ internal class DefaultCpuOpsJvm( @Suppress("UNCHECKED_CAST") override fun matmulWeightTransposed(x: Tensor, weight: Tensor): Tensor { if (weight.shape.rank != 2 || !isHeapPackedWeightForJvm(weight.data)) return super.matmulWeightTransposed(x, weight) + // Already in feed order — the loader produced it that way (#1120). Nothing to permute and + // nothing to cache: the kernels want the other shape label over the same bytes, which costs + // an object rather than a copy of the weight. + rewrapFeedOrderWeight(weight)?.let { return matmul(x, it) } val packed = weight.data as sk.ainet.lang.tensor.storage.PackedBlockStorage val source = packed.packedData val cached = prepackedWeightsJvm.firstOrNull { it.first === source }?.second 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 b84293597..f9f3dabbd 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 @@ -210,12 +210,10 @@ public class StreamingGgufParametersLoader( "ignored; pass only the form." } } - require(form.order == WeightByteOrder.AS_STORED) { - "WeightByteOrder.KERNEL_FEED is not supported by this loader yet (#1120). The bytes are " + - "easy — TensorView.prepack(INPUT_BLOCK_MAJOR) does the permutation — but packed " + - "TensorData addresses packedData as canonical row-major in getBlockScale, getCode " + - "and dequantizeBlock, so feed-order bytes would decode the wrong elements without " + - "failing (#973, #968). Packed storage has to be able to declare its own order first." + require(form.order == WeightByteOrder.AS_STORED || form.shape == WeightShapeOrientation.OUT_IN) { + "WeightByteOrder.KERNEL_FEED needs WeightShapeOrientation.OUT_IN: feed order is defined " + + "relative to a [out, in] weight — which block is 'block b of output row o' has no " + + "answer while the tensor is still labelled in the file's `ne` order." } val requested = form.encoding require(requested !is EncodingRequest.RequantizeTo) { @@ -433,7 +431,42 @@ public class StreamingGgufParametersLoader( "quantizedTensor called for non-quantized type ${tensorInfo.tensorType}" ) } - return ctx.fromData(packed as sk.ainet.lang.tensor.data.TensorData, dtype) + val delivered = if (form.order == WeightByteOrder.KERNEL_FEED) feedOrdered(packed, tensorInfo) else packed + return ctx.fromData(delivered as sk.ainet.lang.tensor.data.TensorData, dtype) + } + + /** + * [packed] with its blocks permuted into the order the packed matmul kernels read, and *saying + * so* (#1120). + * + * The permutation is `TensorView.prepack`, which already owns it and emits the conversion on + * the trace (#1117). What #1120 adds is the second half: the result is rebuilt as a + * `TensorData` carrying `BlockOrder.INPUT_BLOCK_MAJOR`, so every reader is right about the same + * bytes — the kernels address them in feed order deliberately, and anything decoding through + * the view or `toFloatArray()` walks the logical grid and fetches each block from where this + * order put it. Before that, feed-order bytes in a type claiming to be canonical decoded to + * plausible garbage (#1124, #973, #968). + */ + @OptIn(sk.ainet.lang.memory.ExperimentalMemoryApi::class) + private fun feedOrdered( + packed: sk.ainet.lang.tensor.storage.PackedBlockStorage, + tensorInfo: StreamingTensorInfo, + ): sk.ainet.lang.tensor.storage.PackedBlockStorage { + val shape = packed.shape + if (shape.rank != 2 || shape[1] % packed.blockSize != 0) return packed + val prepacked = packed.packedView.prepack(sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR, sink = traceSink) + val bytes = (prepacked.storage as? sk.ainet.lang.memory.Storage.Heap)?.bytes ?: return packed + val order = sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR + return when (tensorInfo.tensorType) { + GGMLQuantizationType.Q4_K -> Q4_KBlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q5_K -> Q5_KBlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q6_K -> Q6_KBlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q8_0 -> Q8_0BlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q4_0 -> Q4_0BlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q5_0 -> Q5_0BlockTensorData(shape, bytes, order) + GGMLQuantizationType.Q5_1 -> Q5_1BlockTensorData(shape, bytes, order) + else -> packed + } } private fun bytesToFloatArray(bytes: ByteArray): FloatArray { diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/KernelFeedOrderTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/KernelFeedOrderTest.kt new file mode 100644 index 000000000..401a47f13 --- /dev/null +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/KernelFeedOrderTest.kt @@ -0,0 +1,163 @@ +package sk.ainet.io.gguf + +import kotlinx.coroutines.runBlocking +import sk.ainet.backend.api.kernel.KernelRegistry +import sk.ainet.context.DirectCpuExecutionContext +import sk.ainet.exec.kernel.ScalarKernelProvider +import sk.ainet.io.JvmRandomAccessSource +import sk.ainet.lang.memory.BlockOrder +import sk.ainet.lang.memory.ExperimentalMemoryApi +import sk.ainet.lang.memory.plan.WeightByteOrder +import sk.ainet.lang.memory.plan.WeightForm +import sk.ainet.lang.memory.plan.WeightShapeOrientation +import sk.ainet.lang.memory.trace.RecordingTraceSink +import sk.ainet.lang.memory.trace.TraceEvent +import sk.ainet.lang.tensor.Shape +import sk.ainet.lang.tensor.Tensor +import sk.ainet.lang.tensor.matmulWeightTransposed +import sk.ainet.lang.tensor.storage.PackedBlockStorage +import sk.ainet.lang.types.FP32 +import java.io.File +import kotlin.math.abs +import kotlin.test.BeforeTest +import kotlin.test.Test +import kotlin.test.assertContentEquals +import kotlin.test.assertEquals +import kotlin.test.assertFailsWith +import kotlin.test.assertTrue + +/** + * #1120: a weight loaded in kernel feed order says so, decodes to the same matrix, and reaches the + * kernels without the per-weight relayout ever running. + * + * Feed order is only meaningful at three-plus blocks per row — at one block per row the two orders + * coincide and every assertion below would hold vacuously (#968). + */ +@OptIn(ExperimentalMemoryApi::class) +class KernelFeedOrderTest { + + private val outDim = 32 + private val inDim = 96 // three Q8_0 blocks per row + + @BeforeTest + fun registerKernels() { + KernelRegistry.register(ScalarKernelProvider) + } + + private fun modelFile(type: GGMLQuantizationType = GGMLQuantizationType.Q8_0): File = + SyntheticGguf.write( + SyntheticGguf.tensor("blk.0.attn_q.weight", type, elements = outDim * inDim) + .copy(dims = listOf(inDim.toLong(), outDim.toLong())), + ) + + private fun load(f: File, form: WeightForm, sink: RecordingTraceSink = RecordingTraceSink()): + Pair, RecordingTraceSink> { + val ctx = DirectCpuExecutionContext() + var w: Tensor? = null + runBlocking { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(f) }, + weightForm = form, + traceSink = sink, + ).load(ctx, FP32::class) { _, t -> w = t } + } + return w!! to sink + } + + private fun asStored() = WeightForm(shape = WeightShapeOrientation.OUT_IN) + private fun feedOrder() = + WeightForm(order = WeightByteOrder.KERNEL_FEED, shape = WeightShapeOrientation.OUT_IN) + + @Test + fun `a feed-order weight declares its order and keeps its shape`() { + val f = modelFile() + try { + val (w, _) = load(f, feedOrder()) + val packed = w.data as PackedBlockStorage + assertEquals(BlockOrder.INPUT_BLOCK_MAJOR, packed.blockOrder, "the bytes must say what order they are in") + assertEquals(Shape(outDim, inDim), w.shape, "feed order is a property of the bytes, not of the shape") + } finally { + f.delete() + } + } + + @Test + fun `feed-order bytes decode to the same matrix they were permuted from`() { + val f = modelFile() + try { + val (canonical, _) = load(f, asStored()) + val (feed, _) = load(f, feedOrder()) + + val canonicalBytes = (canonical.data as PackedBlockStorage).packedData + val feedBytes = (feed.data as PackedBlockStorage).packedData + assertTrue( + !canonicalBytes.contentEquals(feedBytes), + "at three blocks per row the two orders must differ, or this test proves nothing", + ) + + assertContentEquals( + (canonical.data as PackedBlockStorage).toFloatArray(), + (feed.data as PackedBlockStorage).toFloatArray(), + "different bytes, same matrix — that is the whole claim of a declared block order", + ) + } finally { + f.delete() + } + } + + @Test + fun `the product is the same whichever order the weight arrived in`() { + val f = modelFile() + try { + val ctx = DirectCpuExecutionContext() + val xs = FloatArray(inDim) { (it % 13) * 0.0625f } + val x = ctx.fromFloatArray(Shape(1, inDim), FP32::class, xs) + + val (canonical, _) = load(f, asStored()) + val (feed, _) = load(f, feedOrder()) + + val fromCanonical = x.matmulWeightTransposed(canonical).data.copyToFloatArray() + val fromFeed = x.matmulWeightTransposed(feed).data.copyToFloatArray() + + for (o in fromCanonical.indices) { + assertTrue( + abs(fromCanonical[o] - fromFeed[o]) <= 1e-3f * maxOf(1.0f, abs(fromCanonical[o])), + "output[$o]: canonical ${fromCanonical[o]} vs feed-order ${fromFeed[o]}", + ) + } + } finally { + f.delete() + } + } + + @Test + fun `the permutation happens once at load and is reported there`() { + val f = modelFile() + try { + val (_, sink) = load(f, feedOrder()) + val relayouts = sink.events() + .filterIsInstance() + .filter { it.kind.startsWith("prepack") } + assertEquals(1, relayouts.size, "one weight, one permutation: ${sink.events().map { it }}") + + val (_, quiet) = load(f, asStored()) + assertTrue( + quiet.events().filterIsInstance().none { it.kind.startsWith("prepack") }, + "a weight kept as stored is not permuted", + ) + } finally { + f.delete() + } + } + + @Test + fun `feed order without OUT_IN is refused because it would be meaningless`() { + val failure = assertFailsWith { + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(modelFile()) }, + weightForm = WeightForm(order = WeightByteOrder.KERNEL_FEED), + ) + } + assertTrue(failure.message!!.contains("OUT_IN"), failure.message!!) + } +} diff --git a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt index 0a4ac1d08..f79490fbb 100644 --- a/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt +++ b/skainet-io/skainet-io-gguf/src/jvmTest/kotlin/sk/ainet/io/gguf/WeightFormLoaderParityTest.kt @@ -140,18 +140,26 @@ class WeightFormLoaderParityTest { } @Test - fun `KERNEL_FEED is refused for the reason it is refused`() { - // Not "unsupported": the bytes are the easy part. The refusal is because packed TensorData - // reads packedData as canonical row-major, so feed-order bytes decode wrong silently (#1120). + fun `KERNEL_FEED is accepted now that packed storage can declare its order`() { + // This slice refused it: packed TensorData read packedData as canonical row-major, so + // feed-order bytes decoded to the wrong elements silently. #1120 gave the bytes a way to + // say what order they are in, so the refusal is gone — replaced by the one constraint that + // remains real, since feed order is defined relative to an [out, in] weight. + StreamingGgufParametersLoader( + sourceProvider = { JvmRandomAccessSource.open(file()) }, + weightForm = WeightForm( + order = WeightByteOrder.KERNEL_FEED, + shape = WeightShapeOrientation.OUT_IN, + ), + ) + val failure = assertFailsWith { StreamingGgufParametersLoader( sourceProvider = { JvmRandomAccessSource.open(file()) }, weightForm = WeightForm(order = WeightByteOrder.KERNEL_FEED), ) } - val message = failure.message!! - assertTrue(message.contains("#1120"), "it names where this is being fixed: $message") - assertTrue(message.contains("canonical row-major"), "and why it cannot be faked: $message") + assertTrue(failure.message!!.contains("OUT_IN"), failure.message!!) } @Test diff --git a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt index 6cefa7266..868487084 100644 --- a/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt +++ b/skainet-lang/skainet-lang-core/src/commonMain/kotlin/sk/ainet/lang/tensor/storage/PackedBlockStorage.kt @@ -93,6 +93,21 @@ public interface PackedBlockStorage { public fun toFloatArray(): FloatArray { val result = FloatArray(shape.volume) + if (blockOrder == sk.ainet.lang.memory.BlockOrder.INPUT_BLOCK_MAJOR && shape.rank == 2) { + // Feed-order bytes hold block `(o, b)` at physical index `b * rows + o`, so decoding + // them in physical sequence would emit the matrix transposed-in-blocks. Walk the + // *logical* grid instead and fetch each block from where this order put it — which is + // what makes a feed-order weight decode to the same matrix as the canonical one it was + // permuted from (#1120). + val rows = shape[0] + val blocksPerRow = shape[1] / blockSize + for (o in 0 until rows) { + for (b in 0 until blocksPerRow) { + dequantizeBlock(b * rows + o, result, (o * blocksPerRow + b) * blockSize) + } + } + return result + } var offset = 0 for (i in 0 until blockCount) { val remaining = shape.volume - offset