Skip to content

One tensor DSL whatever the weight is packed as: w.t() must not throw #1108

Closed

Description

@michalharakal

Follows #973 / #1096.

The problem

After #1096 a model author writes different code depending on how the weight is stored:

x.matmul(w.t())                 // dense — fine
x.matmulWeightTransposed(w)     // block-quantized — because t() throws

That is a leak. Whether a weight arrives packed depends on the file the user loaded and the target it is running on — neither of which the model author knows or should encode. The same nn { } graph has to work over an FP32 safetensors checkpoint on a desktop and a Q4_K GGUF on a phone. Right now the second one throws from ops.transpose.

Linear hides this today because it calls the primitive itself. Anything a user hand-writes does not.

Why it was banned, and why a fix is not simply "un-ban it"

transpose on packed data used to be a metadata relabel over the same bytes — the shape changed and the payload did not. NativeLazyTransposeGroundTruthReproTest demonstrates the failure precisely: canonical row-major packing and input-block-major packing of the same logical matrix are bytewise different, and a bare relabel hands the kernel the wrong one. It does not fail; it returns plausible garbage (#968, SKaiNET-transformers#307).

So the fix cannot be "make t() return a relabelled tensor again". It has to be a value that knows it is not materialized.

Proposal

transpose(w) on a packed weight returns a transposed-weight view: a TensorData wrapper meaning "logically Wᵀ; the bytes are still W's packing; nothing may read them as if they were Wᵀ's".

  • matmul(x, <that view>) recognises it and routes to matmulWeightTransposed(x, w) — the relayout runs once and is cached, exactly as now. This is the eager flattening of matmul(x, transpose(w)), achieved by a marker rather than by an optimizer pass, because eagerly there is no pass to run.
  • t().t() returns the original weight — transpose becomes an involution again, which it has not been since packed data existed.
  • Anything that wants the transposed values — get(), an elementwise op, materialize() — either requantizes explicitly at that point, or throws there. That is where the operation is genuinely undefined; t() is not.

The safety property to enforce and test: no consumer can reach packedData through the marker. That is the whole difference between this and the bug above, so the tests must be adversarial about it — including a fixture at least three blocks wide, since at one block per row the two orderings coincide and any such test passes vacuously.

Acceptance

  • x.matmul(w.t()) produces bit-identical output to x.matmulWeightTransposed(w) for every packed encoding, on a ≥3-block-wide fixture
  • w.t().t() === w
  • Reading values through the marker either requantizes correctly or fails loudly; a test proves the packed bytes cannot be consumed raw through it
  • The relayout still happens once per weight, not once per call — the [973.3] A weight-transposing matmul primitive; ops.transpose on packed data becomes a loud error #1096 allocation test still holds
  • NativeLazyTransposeGroundTruthReproTest stays green

Related: #1096 added the fluent Tensor.matmulWeightTransposed(w) extension, which makes the explicit form pleasant. This issue is about not needing the explicit form.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    architecturaldsl/apiDSL and API designtensorsTensor operations and data structures

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions