When Weights Move to the Left: A Transpose Identity That Moves the Whole Compiler Chain

Series contents · Operators and Layouts · 阅读中文版

“It is still matrix multiplication. Why can it no longer compile just because the weights are on the left?”

From linear algebra’s perspective, that is a reasonable question. From the backend’s perspective, matrix-unit dataflow, weight packing, bias parameters, and quantization channels may all be organized around dynamic input @ constant weights. Move the constant to the left, and the operator’s name stays the same while almost every implementation assumption is touched.

The Transpose Identity and Shape Proof

The most direct—and dangerous—approach is to swap inputs. Matrix multiplication is generally not commutative; even the output shape can change. The correct teaching route is:

1
2
3

Y = W @ A
Y^T = A^T @ W^T
Y = (A^T @ W^T)^T

Constant W can be transposed at compile time. Transposing runtime A and restoring the result require corresponding graph nodes. The core computation then returns to the existing path with constant weights on the right.

Verify it with small matrices instead of memorizing a formula:

1
2
3
4
5
6
7
8

W = [[ 1, 2, 0],
     [-1, 0, 3]]
A = [[2, 1],
     [1, 4],
     [3,-2]]

W@A = [[4, 9],
       [7,-7]]

Computing A^T@W^T gives [[4,7],[9,-7]]; transposing restores the result exactly. This checks dimensions as well as values: W is M×K and A is K×N. The intermediate multiplication is N×K by K×M, producing N×M, which is finally restored to M×N.

Batch Dimensions and Preconditions

With batch dimensions, transposing the entire tensor indiscriminately will not do. Suppose A:[B0,B1,K,N] shares W across batches. Flatten the batch dimensions to B=B0*B1, yielding [B,K,N]; swap the last two dimensions to obtain [B,N,K]; multiply by transposed W to produce [B,N,M]; swap back to [B,M,N]; then restore [B0,B1,M,N].

The inspected historical lowering used this shape chain, creating reshape and permute operations where needed. It did not support arbitrary MatMul unconditionally: the left operand had to be constant, the right operand could not be a weight, the original operation could have no bias, and existing input or output transpose flags prevented entry into this dedicated pattern. Constant shapes and the dynamic operand’s rank also had restrictions. A public account should call this a bounded support extension, not a general matrix transformer.

Why initially exclude bias? In Y=W@A+b, bias might broadcast along output columns or have another shape. Transposing changes its broadcast axis. Feeding the original bias unchanged into the new MatMul can turn “add by column” into addition along another direction. It is better to reject the case initially and design a separately provable bias mapping than to be silently wrong in the first version.

Quantization Roles and Bias Compensation

Quantization raises the problem another level. A right-weight path typically prepares scales and bias around output channels. Before transposition, W’s output rows correspond to the original output’s M dimension. After transposition, they become columns of W^T and the new computation’s output channels. If per-channel scales were bound to another axis, copying an array unchanged cannot make the meanings equivalent.

Historical code created a new constant for the transposed weights and then read its quantization information. Inspecting helper interfaces showed that weight scales were looked up by weight name in an external mapping. The new constant’s name, derived metadata, and channel semantics therefore had to connect; transposed numerical values alone were insufficient. The same commit retained parameter encodings not immediately matched to weight nodes, preserving information for later generation or lookup. That does not automatically prove every new name and per-channel axis correct.

Stronger regressions should use different scales per channel and check the semantic coordinates to which scales bind after transposition, rather than only array length.

Import also contains an operand-numbering trap. If activation quantization has just one encoding and the model places weights at operand 0 and dynamic activation at operand 1, the usual assumption that “the first input encoding belongs to the first operand” fails. A historical change recognized this static-left MatMul case and skipped activation-encoding matching for the left weight, letting the dynamic side receive the correct type. Dataflow roles are more reliable than position numbers.

Bias parameters may be nonempty even when the mathematical model has no bias. Quantized integer multiply-accumulate needs compensation for the input zero point. In a simplified case, if dynamic input zero point is z and quantized weights for one output channel are q_w, expanding the multiplication introduces a constant compensation related to -z*sum(q_w). The historical change extracted bias calculation from an interface tied directly to the old operation object into one accepting input, weight, output shape, and transpose conditions, allowing the new path to call it.

The general lesson is to make hidden assumptions explicit before reusing code: which operand is the input, which is the weight, and where are the output channels? Otherwise, a function still called “compute MatMul bias” may keep looking at the original right operand, leaving the transpose transformation only half implemented.

Compilation Costs and Correctness Tests

Compilation performance also has a bill. Transposing a constant allocates a new array and may increase peak compilation memory and artifact size. Runtime transposes of A and the result may add two data movements. If neighboring operators can directly produce or consume the required layout, some transposes can be absorbed. Otherwise, broader compilability does not imply ideal speed. Capability support and performance optimization are different promises.

A comparison experiment could hold multiply-accumulate workload fixed while varying M, N, K, and the number of batch dimensions, recording constant preprocessing time, temporary-buffer sizes, transpose traffic, and core MatMul time separately. Pay particular attention to tiny matrices with many batches: launches and transposes may stand out more than multiply-accumulates. The three transpose marks in the identity are not three free mathematical decorations.

Correctness tests can begin with the nonsquare matrices above, avoiding square shapes that let dimensional errors slip through. Add two batch dimensions, dynamic inputs with distinct labels per batch, a nonzero input zero point, different per-channel scales, and negative cases for restricted bias and transpose attributes. After restoring shape, check every batch instead of comparing only aggregate statistics after flattening.

Failure propagation needs verification too. Failure in derived-weight creation, quantization lookup, compensation-parameter construction, or weight lowering must not leave an apparently successful new MatMul. Completing the shape and metadata plan before changing IR can reduce incomplete graphs after a mid-rewrite failure.

Quantization Axes, Sharing, and Boundaries

To make quantization axes concrete, bind a different scale to each row of W. After transposition, those scales should follow the original row elements and become interpretation rules for corresponding columns of W^T. Scale values may stay unchanged while their axis position changes. If the original scales actually apply along K, transposition places them in another semantic position; they cannot simply be called the new output-channel scales. Migration rules must follow the original encoding’s meaning, not the observation that some dimension happens to have the right length after transposition.

A well-designed derived-constant interface could therefore return more than a new tensor: it could include an old-to-new coordinate mapping and transformed, or pending-transformation, quantization descriptions. Later bias and packing stages would no longer need to guess lineage from names. The inspected historical implementation mainly connected new weights through external encoding lookup. The more explicit interface proposed here illustrates how coupling could be reduced, not a refactor already completed.

Flattening batches has prerequisites as well. If W is shared across batches, flattening all dynamic batch dimensions preserves weight selection. If left-hand weights themselves contain different data per batch, their leading dimensions are not mere decoration. The historical dedicated path restricted left-weight shapes to controlled leading unit dimensions, avoiding a promise of full batch broadcasting. Further extension should first define the broadcast result of left and right batch dimensions, then consider whether flattening preserves correspondence.

A teaching test can give every batch a unique marker, for example by adding different small integer offsets to one base matrix. After restoring batch shape, verify that each result retains its marker. Flattening the entire output and comparing its sum may miss a swapped batch order. Identical batch inputs would defeat this test too, so deliberately avoid symmetric data.

Compile-time weight transposition can also be repeated unnecessarily. If several static-left MatMuls share a constant and each lowering creates a separate transpose, memory use and artifact size may grow. Caching or common-subexpression reuse could help, but the cache key should include content, element type, quantization encoding, and target storage format. Caching only by the original weight name risks wrong reuse for duplicate names or different quantized views.

Finally, batch-dimension products need checks for overflow and dynamic values. The mathematical expression B=B0·B1 looks natural, but host integers have limits and dynamic dimensions may use special sentinel values. A path called “static left” does not imply that every dimension on the right is static and valid. Review should locate those prerequisites layer by layer. If they cannot be found, record them as constraints requiring verification rather than supplying guarantees from a function’s name. This honesty about boundaries is worth preserving when turning real engineering experience into public technical writing.

An elegant identity is only the key to the door. Passing through it requires shapes, constant bytes, quantization roles, bias compensation, and runtime layouts to move together. Mathematics proves that a new arrangement solves the same problem; the compiler must ensure that every layer actually adopts that arrangement.

Series contents · Previous

• ← Previous Post
• Next Post →