Skip to content

fix(cpu): stop relayouting a packed weight for kernels this backend does not have - #1126

Merged
michalharakal merged 1 commit into
developfrom
fix/1124-packed-block-order
Aug 25, 2026
Merged

michalharakal merged 1 commit into
developfrom
fix/1124-packed-block-order

Conversation

@michalharakal

Copy link
Copy Markdown
Contributor

Closes #1124.

The common DefaultCpuOps computed packed matmul wrongly whenever the KernelRegistry was empty — silently, by several times the magnitude of the answer:

[32 x  96] expected 8.25   got -3.0625
[64 x  96] expected 8.25   got  9.9375
[32 x 128] expected 12.625 got -3.5

Ground truth is the weight's own toFloatArray(), so the packed path was disagreeing with its own decoder.

What was actually wrong

matmulWeightTransposed relayouted the weight into the order the JVM vectorized kernels read, then handed the result to whatever computed the product.

On the JVM that is chooseQuantizedMatmul, which addresses packedData in feed order deliberately, and is correct. The common implementation has no such kernel — it decodes through views. And what it was decoding is a tensor whose shape says [in, out] while its blocks still run along the original input dimension. That is neither canonical nor input-block-major-at-that-shape: it is a private artifact that only its producer understands, and reading it as either gives plausible garbage.

So relayouting for a kernel that does not exist was pure harm.

The fix

The common path decodes the weight where it lies. KernelDispatch.matmul already wants the weight output-major, which is exactly the shape an [out, in] weight has, so the canonical view goes straight in with nothing rearranged.

DefaultCpuOpsJvm overrides matmulWeightTransposed to keep the relayout-and-cache path — there the consuming kernels exist and #1096's once-per-weight conversion still pays. The split now matches reality: fast bytes where something reads them, correct decoding where nothing does.

Block order reaches views now

PackedBlockStorage.blockOrder, defaulted to ROW_MAJOR so nothing existing changes, forwarded by packedView into Layout.blocked — which has accepted a blockOrder since #1094 and never received one. That closes the information gap #1120 needs.

Worth recording what it cannot do, since #1120 will build on it: the relayout artifact above has no honest block order to declare. Marking it INPUT_BLOCK_MAJOR was the first fix I tried, and it produced zeros — which is what showed the shape and the blocks were inconsistent in the first place.

Why nothing caught this

  • skainet-backend-cpu's JVM tests always resolve to DefaultCpuOpsJvm, whose override is correct.
  • Its native / JS / Wasm tests do use the common implementation, but their platform factories register ScalarKernelProvider first.

So every packed-matmul correctness test in the tree — including the ones added for #973, #1096 and #1108 — validated a configuration that was not the broken one. The gap was structural, not a missed case in any single test.

PackedMatmulEmptyRegistryTest pins the uncovered configuration: four encodings × three shapes, three blocks per row so canonical and feed order are distinguishable (#968), asserted against the weight's own decoder. It lives in commonTest, so it runs on the targets where that path is real.

Impact

Production paths register a provider, so applications were not hitting this. What was exposed: anything constructing ops without PlatformCpuOpsFactory, any target where registration fails — and the fallback's whole purpose, which is to work when nothing else is available.

Gate

scripts/pr-gate.sh — all legs passed.

🤖 Generated with Claude Code

…oes not have

Closes #1124.

The common DefaultCpuOps computed packed matmul wrongly whenever the
KernelRegistry was empty — silently, by several times the magnitude of the
answer. Ground truth is the weight's own toFloatArray(), so the packed path
was disagreeing with its own decoder.

matmulWeightTransposed relayouted the weight into the order the *JVM*
vectorized kernels read, then handed the result to whatever computed the
product. On the JVM that is chooseQuantizedMatmul, which addresses
packedData in feed order deliberately and is correct. The common
implementation has no such kernel: it decodes through views, and what it was
decoding was a tensor whose shape says [in, out] while its blocks still run
along the original input dimension. That combination is not canonical and is
not input-block-major-at-that-shape; it is a private artifact that only its
producer understands, and reading it as either gives plausible garbage.

So the common path no longer relayouts. KernelDispatch.matmul already wants
the weight output-major, which is exactly the shape an [out, in] weight has,
so the canonical view goes in with nothing rearranged. DefaultCpuOpsJvm
overrides matmulWeightTransposed to keep the relayout-and-cache path, where
the kernels that consume those bytes actually exist and #1096's once-per-weight
conversion still pays.

Also plumbs block order through to views: PackedBlockStorage.blockOrder,
defaulted to ROW_MAJOR so nothing existing changes, forwarded by packedView
into Layout.blocked — which has accepted a blockOrder since #1094 and never
received one. That closes the information gap #1120 needs. It is worth
recording what it cannot do: the relayout artifact above has no honest block
order to declare, so marking it was the first fix I tried and it produced
zeros.

Why no test caught the original bug: skainet-backend-cpu's JVM tests always
resolve to DefaultCpuOpsJvm, and its native/JS/Wasm tests do use the common
implementation but their platform factories register ScalarKernelProvider
first. Every packed-matmul test in the tree, including those added for #973,
#1096 and #1108, validated a configuration that was not the broken one. The
new test pins the empty-registry configuration explicitly, in commonTest so
it runs where that path is real.

Gate: scripts/pr-gate.sh — all legs passed.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
@github-actions

Copy link
Copy Markdown

📖 Documentation Preview

The documentation has been built successfully for this PR.

Generated Files:

  • Operator documentation: docs/modules/operators/_generated_/
  • JSON schema output: operators.json

Artifacts:

  • Download the documentation-preview-1126 artifact to view the complete documentation locally.

This comment will be updated automatically when the PR is updated.

@michalharakal
michalharakal merged commit 4b7497b into develop Aug 25, 2026
20 checks passed
@michalharakal
michalharakal deleted the fix/1124-packed-block-order branch August 25, 2026 15:44
yasanglass pushed a commit to yasanglass/skainet that referenced this pull request Oct 21, 2026
Closes SKaiNET-developers#1120. Follows SKaiNET-developers#1124/SKaiNET-developers#1126.

A weight can now be permuted into the order the packed matmul kernels read at
load time, and the resulting TensorData declares BlockOrder.INPUT_BLOCK_MAJOR
rather than pretending to be canonical. Every reader is then right about the
same bytes: the kernels address them in feed order deliberately, while
toFloatArray() and view decoding walk the logical grid and fetch each block
from where that order put it.

The payoff is that SKaiNET-developers#1096's O(bytes) per-weight relayout stops running at all
for such a weight. DefaultCpuOpsJvm relabels the shape over the same array
instead — asserted structurally with assertSame on the byte array, since a
relayout would necessarily produce a different one, and timing proves nothing.

Two things that had been conflated are now separate.
relayoutPackedWeightForKernels still produces bytes labelled [in, out] for the
JVM kernels; that is a private artifact whose shape and blocks disagree, which
is exactly why SKaiNET-developers#1126 found it undecodable and why it keeps no declared order.
A loader-produced feed-order weight is a different thing: it keeps its
[out, in] shape and describes itself.

The permutation reuses TensorView.prepack(INPUT_BLOCK_MAJOR), which already
owned it and already emits the conversion on the trace — so SKaiNET-developers#1117 reports the
prepack without new code, and a test counts exactly one per weight and none
for a weight kept as stored.

KERNEL_FEED requires WeightShapeOrientation.OUT_IN, and says so: which block
is "block b of output row o" has no answer while the tensor is still labelled
in the file's ne order, so refusing beats permuting against a meaningless grid.

SKaiNET-developers#1115's test asserting the old refusal is rewritten rather than deleted — it
now pins the constraint that remains.

Gate: scripts/pr-gate.sh — all legs passed. One earlier run failed
SlicingTest.testPerformanceAccessPatterns on ChromeHeadless, an unrelated
dense-slicing test that passes on its own; not touched here.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Packed matmul returns wrong numbers when DefaultCpuOps runs with an empty KernelRegistry

1 participant