Fewer Operations—So Why Disable the Optimization?

Series contents · Resources and Scheduling · 阅读中文版

“Instead of multiplying the matrices and then taking a small slice, let’s slice the weights first and do a smaller multiplication.”

“The equation is fine. Which quantization parameters belong to the new weights?”

“Don’t the parameters follow the name?”

Algebraic transformations look clean on a whiteboard. A value inside a compiler, however, carries more than a shape and elements. It also carries scales, zero points, a channel axis, layout, storage, and provenance. If an optimization leaves that information behind while saving multiplications, it becomes a way to make errors harder to notice.

Start with a valid floating-point identity

Let X have shape B×K and W have shape K×N, and suppose only certain output columns are needed:

1

slice_columns(X * W) = X * slice_columns(W)

Under suitable slicing conditions, this transformation can reduce computation or intermediate output. If several consumers take different columns, it may also split one large matrix multiplication into several smaller ones.

The equation describes numerical operations, but says nothing about how quantization parameters change. Suppose W uses per-output-channel quantization, with a separate scale and zero_point for each column. Selecting a set of columns J must select the parameters using the same J. If a newly named weight object copies only bytes and shape without registering its derived quantization information, downstream code may fail to find the parameters or use scales belonging to different columns.

Such defects can easily remain invisible in floating-point tests. A floating-point tensor does not have that additional chain of per-channel parameters, so the algebraic identity holds without it.

The verified action is more conservative than the comment might suggest

The relevant commit adds an explanation near the rewrite logic: splitting quantized weights produces weight objects with new names, while the derived per-channel quantization information is not yet propagated safely. It also removes the entire “matrix multiplication followed by slicing” pattern from the canonicalization-pattern registration list.

That distinction matters. The code contains a draft condition targeting quantization information, but it is commented out. The action that actually takes effect disables the entire pattern. It does not merely skip quantized weights while continuing to optimize floating-point ones.

A public account should therefore not describe this as a finely targeted, quantization-aware rewrite. The supported conclusion is that the original graph is retained until the full path for derived metadata is in place. This is a broad correctness safeguard that may give up some otherwise safe optimization opportunities.

The commit adds no dedicated regression test. We can confirm the change to pattern registration; the change alone does not establish numerical validation for all affected models.

Why looking up parameters by global name makes rewriting harder

If quantization information is retrieved from an external table by name, producing a new value involves more than creating an IR node. A rewrite must also maintain the relationship between names and metadata. Copying, slicing, concatenation, and transposition each impose different requirements on that relationship.

One general direction is to place semantically essential parameters in types or explicit attributes where possible, requiring rewrite functions to handle them. Even when external data remains necessary, a common API for deriving weights could commit the value, quantization parameters, and provenance mapping together, instead of leaving each pattern to patch names independently.

Moving parameters into types is no universal solution: large arrays affect IR size, serialization, and comparison costs, while transformations involving the channel axis still require correct reasoning.

A different optimization: avoid filling the tensor with broadcast constants

Another change in the same history addresses a different issue. Scale and bias parameters were previously expanded to match the input shape. The change retains channel vectors and constructs compact broadcast shapes instead.

Consider a teaching input of shape B×T×C. If each channel’s coefficient is identical across batches and positions, expanding scale from C elements to B×T×C elements remains mathematically correct, but materializes repeated information. For the teaching values B=2 and T=37, the element counts differ by a factor of 74. That is a storage ratio derived from the shapes, not a measured speedup.

A compact representation can reduce the potential cost of generating, storing, and transferring constants. It relies on a clear premise, however: downstream operations must understand that broadcast shape, and the channel axis must already have been normalized to the agreed position. If the input layout still permits several interpretations, simply placing channels in the final dimension can turn “save memory” into “scale along the wrong axis.”

Both changes ask the same question: who carries the semantics?

Weight slicing goes wrong when it drops metadata that must change with the value. Parameter expansion is inefficient when it repeatedly stores information that broadcasting rules could express. One needs to preserve more semantic information; the other needs to remove redundant representation.

“The simpler the IR, the better” is therefore incomplete. A small node can conceal an unverified broadcasting assumption. A seemingly verbose matrix multiplication followed by a slice may preserve the quantization relationships the current system actually understands.

Begin an optimization review by listing invariants: which dimensions are selected, which channel parameters must be reordered, how the output type changes, whether layout changes, whether the numerical range changes, and whether all derived weights remain discoverable. Only after answering these questions should the discussion turn to multiplication counts and bytes.

One concrete counterexample can explain more than ten random models

Take teaching weights with three columns whose scales are s0, s1, and s2, all clearly different. Select only the last two columns. Their correct new parameters are s1 and s2. If the rewrite uses s0 and s1 instead, the output shape is still correct, but the values are scaled by the wrong factors.

Setting every scale to the same number hides the effect of that incorrect mapping. Test data should deliberately break symmetry: different scales and zero points across channels, a slice starting at a nonzero position, and uneven partitions. Such a counterexample explains a break in metadata propagation more clearly than a large random graph.

Broadcasting tests should likewise avoid equal lengths on every axis. If B, T, and C happen to be equal, broadcasting along the wrong axis may still pass shape checks. Distinct small dimensions make the semantics easier to observe.

When can the optimization be enabled again?

At least four conditions should hold: derived weight metadata has a reliable source; the relationship between the slice axis and quantization axis is handled explicitly; downstream lowering can consume the new parameters; and regression tests cover deliberately asymmetric counterexamples.

Then choose a policy. An initial implementation might allow only per-tensor quantization, or particular transformations along an axis other than the quantization axis. Alternatively, it could implement full parameter propagation for per-channel slicing. Either choice must be demonstrated by active code and tests. An ideal condition in a comment is not current behavior.

Any performance gain must also be weighed against more small matrix tasks, repeated reads of the input, kernel launches, and synchronization. Fewer multiplications do not guarantee lower end-to-end latency, especially if splitting the work reduces the utilization achieved by a large matrix multiplication.

Make a ledger for derived metadata

Suppose a rewrite creates two new weights. For each new value, a reviewer can make a ledger: which range supplies the original elements, how the shape changes, whether the quantization axis moves, where the scales and zero points come from, how debug provenance maps back, and which identifier later serialization uses. An unanswered item indicates that the rewrite’s semantics are still incomplete.

This ledger also distinguishes “propagation is impossible” from “propagation is not implemented yet.” The first may mean that the transformation is invalid under a particular quantization representation. The second may be resolved by adding an explicit data transformation. Labeling both simply “unsupported” leaves future maintainers unsure whether to add an algorithm or tighten the applicability conditions.

Broadcast constants deserve the same ledger. State which axis the compact scale vector describes, whether the bias quantization scale relates to the input or the output, and whether expanding a scalar parameter into a channel vector preserves its numerical meaning. All-zero constants can create a false sense of safety because zero seems to broadcast correctly everywhere. Tests also need nonzero, asymmetric parameters.

Finally, examine the computation graph. If several slices collectively cover almost every output column, splitting may save little arithmetic while increasing scheduling work and input reads. If only a small fraction is needed, the optimization opportunity is clearer. Applicability conditions, a semantic ledger, and a cost model are three separate requirements for re-enabling the pattern. Calling a patch an “optimization” supplies none of them.

Regression matrix and experiments not yet run

A proposed first step is to check quantized-slicing equivalence in an independent, low-dimensional integer reference program, then compare compact and expanded parameters element by element. Performance experiments should separately measure compilation time, constant size, runtime transfers, and computation time, avoiding claims of overall acceleration based on a single declining metric.

Sometimes the most responsible optimization patch first disables an optimization. What deserves to be restored is more than the whiteboard equation: it is the complete set of verifiable semantics behind it.

Back to series contents · Previous

• ← Previous Post
• Next Post →