Packed-quant lazy transpose silently corrupts matmul for any weight with >1 block/row (all formats, not just Q5_0/Q5_1) #968

Summary

ops.matmul(x, ops.transpose(W)) — the "classic" lazy-transpose dispatch path every Linear-style layer uses when it hasn't been told to pre-transpose its packed weights — silently produces wrong output for any packed/quantized weight served by the packed-quant matmul dispatch (DefaultCpuOps.chooseQuantizedMatmulHeap, native/Panama/scalar kernel tiers alike) whenever the weight has more than one quantization block per row (inputDim > blockSize, i.e. essentially every real model — hidden dims are almost always many multiples of the 32/256-element block size).

This is general to every packed format the matmul dispatch serves — Q4_0, Q5_0, Q5_1, Q8_0, Q4_K, Q5_K, Q6_K — not Q5-specific, despite surfacing first via Q5_0/Q5_1 downstream. It reproduces on the scalar and Panama-vector kernel tiers exactly like the native (FFM) tier; what changes with the native tier active is only that a specific downstream test's random weight data happened to manifest as all-zero output rather than "merely" wrong.

Discovered downstream in SKaiNET-transformers#307, which observed: constructing Q5_0BlockTensorData/Q5_1BlockTensorData directly and driving them through linearProject's classic lazy-ops.transpose branch — with skainet-backend-native-cpu on the classpath so the native kernel is selected — produced all-zero matmul output, while the same weights through a pre-transposed path (skipping the lazy transpose, constructing the weight already in kernel-native layout) matched the FP32 reference byte-for-byte. PR #307's own workaround was to default the GGUF→packed-weight converters to the pre-transposed layout, correctly noting in its description: "this reproduces identically regardless of the pre-transposed default flip, so it's a pre-existing engine 0.40.0 issue... not something introduced or fixable in [the transformers] repo."

That workaround is not sufficient by itself. Any code path that still constructs a canonically-packed (row-major) weight and drives it through ops.transpose + the packed-quant matmul dispatch is silently wrong — not loudly erroring, returning zeros or plausible-looking garbage instead. That is the more dangerous failure mode: a future caller, a different downstream consumer, or a code path the pre-transposed default doesn't reach (PR #307 itself notes Q4_0/Q5_0/Q5_1 remain conditionally gated, not unconditionally pre-transposed) can reintroduce this silently.

Reproduction (engine-native, no transformers dependency)

Added NativeLazyTransposeGroundTruthReproTest (skainet-backends/skainet-backend-native-cpu/src/jvmTest/kotlin/sk/ainet/exec/kernel/), one case per packed format the native provider serves. Methodology: build the SAME logical weight matrix as two bytewise-different packings —

• canonical (row-major): for output row o, its blocksPerInputDim blocks are contiguous, then row o + 1 — what a [outputDim, inputDim]-shaped weight naturally has when loaded verbatim (e.g. from GGUF) or built any other row-major way.
• kernel-native (input-block-major): (blockIdx * outputDim + o) — every packed-quant native/Panama/scalar matmul kernel's actual physical layout assumption (see the Per-block packed weight layout comment in q5_0_matmul.c / q4_0_matmul.c / q4k_matmul.c / etc.).

Then, for each format:

• classic path: construct the weight canonically with shape [outputDim, inputDim], call ops.transpose, ops.matmul(x, transposed).
• pre-transposed path: construct the SAME logical weight already in kernel-native bytes with shape [inputDim, outputDim], skip ops.transpose, ops.matmul(x, w) directly.
• ground truth: independently computed via PackedBlockStorage.toFloatArray()'s block-sequential dequant of the canonical bytes (reshaped [outputDim][inputDim]) — a different code path from the matmul kernel under test.

Before the fix, on develop (native provider forced via KernelRegistry.register(NativeKernelProvider)):

Every format is wrong whenever blocksPerInputDim > 1; the single-block-per-row control (inputDim == blockSize) passes, which is exactly why some existing tests never caught this — several used degenerate single-block dimensions. On synthetic random weight bytes the "classic" output isn't literally all zero (it's wrong-but-nonzero); the downstream all-zero symptom is a property of the specific model weight byte patterns PR #307's test happened to exercise, not evidence the bug is narrower than described here.

Root cause

DefaultCpuOps.transpose() (skainet-backends/skainet-backend-cpu/src/commonMain/kotlin/sk/ainet/exec/tensor/ops/DefaultCpuOps.kt) implements the packed-quant "lazy transpose" as a bare shape relabel: given a Q5_0BlockTensorData/Q4_KBlockTensorData/etc. with shape (rows, cols), it returns a new instance with shape (cols, rows) wrapping the exact same packedData byte array, on the documented claim that "the matmul kernels index the packed bytes input-block-major from the post-swap shape, so transpose is a pure shape swap — same bytes, no copy."

That claim is true only when the bytes were already physically arranged input-block-major ((blockIdx * outputDim + o), i.e. kernel-native) before the swap — a shape relabel cannot reorder bytes. But a freshly-constructed or GGUF-loaded packed tensor's bytes are canonical / row-major ((o * blocksPerInputDim + blockIdx), output-row-outer): this is what PackedBlockStorage.toFloatArray()'s block-sequential dequant assumes, and how any straightforward row-major producer (including the engine's own StreamingGgufParametersLoader.quantizedTensor → fromRawBytes) hands off packed bytes.

Canonical order and kernel-native order are literal transposes of the (outputDim, blocksPerInputDim) block grid (treating each block as an atomic item), and they coincide only when blocksPerInputDim == 1 — a single block per row. For every wider row, the shape-relabel-only "transpose" hands the matmul kernel bytes in the wrong physical order, and the kernel — having no way to detect this — reads garbage scale/code fields and produces silently wrong (sometimes all-zero, sometimes merely incorrect) output. DefaultCpuOpsJvm.transpose() had an additional, independently-buggy duplicate of the same shape-relabel-only logic specifically for the Q4_KTensorData (ByteArray-backed) case.

This is a general dispatch-contract bug, not a Q5_0/Q5_1-specific one: every packed format's transpose() arm made the identical incorrect assumption, and every native/scalar/Panama matmul kernel makes the identical fixed kernel-native layout assumption regardless of how its input tensor's shape was obtained.

Fix

Filed alongside this issue: [link to be added — PR branch fix/q5-native-lazy-transpose-allzero].

DefaultCpuOps.transpose() now performs the real O(bytes) block-grid permutation (transposePackedBlocks) — physically reordering canonical bytes into kernel-native layout — instead of a bare shape swap, for all seven packed formats (Q4_0, Q5_0, Q5_1, Q8_0, Q4_K, Q5_K, Q6_K). This is strictly more expensive than the old "free" relabel (a real memcpy-style permutation instead of O(1)), but still far cheaper than the FP32 dequant round-trip the lazy-transpose optimization exists to avoid, and — unlike the old version — it is correct. DefaultCpuOpsJvm's redundant, independently-buggy Q4_KTensorData interception was removed in favor of falling through to the shared, now-corrected base implementation.

A misaligned packed tensor (inputDim not a multiple of the format's block size) now fails loudly (IllegalArgumentException) rather than silently truncating a partial block during the permutation.

Test plan

• New NativeLazyTransposeGroundTruthReproTest (native-cpu jvmTest): one ground-truth-vs-classic-vs-pre-transposed case per packed format, all asserting non-zero, matching output — this is the regression guard for this issue.
• PackedMatmulDispatchTest (backend-cpu commonTest/jvmTest/linuxX64Test) updated to build genuinely canonical (row-major) synthetic bytes instead of already-kernel-native bytes, so it actually exercises the real ops.matmul(x, ops.transpose(W)) contract instead of self-fulfilling around the bug.
• QuantizedMemSegMatmulTest (backend-cpu jvmTest) updated: the Q4_K/Q6_K "lazy transpose keeps the same packedData reference (zero-copy)" assertions are now content-equality checks, since zero-copy is no longer the correctness contract.
• Full skainet-backend-cpu and skainet-backend-native-cpu module test suites green on jvmTest and linuxX64Test; apiCheck green (no public API signature changes).

🤖 Generated with Claude Code