Triton PR #11613: Why Transposing an FP4 Dot Doubled the Wrong Dimension

Triton PR #11613 fixes a compiler transformation that preserved the matrices, scales, and element types of a scaled dot product but lost two packing flags.

That omission changed the logical shape of an FP4 operand.

The original kernel stored two e2m1 values in each byte along the output-column dimension. After AccelerateMatmul transposed the dot product, that same physical data needed to be described as two values packed along the row dimension.

The transformation swapped the operands but allowed both packing flags to fall back to their default value of true. Triton then interpreted the packed byte as extending the reduction dimension instead.

In the reported case, the logical K dimension changed from 32 to 64. Scale validation consequently derived a group size of 64, rejected it, and terminated the compiler pass with:

LLVM ERROR: scale factor must be 16 or 32. Got 64

The fix is three lines: swap the packing flags when the operands are swapped.

FP4 needs more than a dtype and a physical shape

Triton’s tl.dot_scaled supports microscaling formats such as e2m1, an FP4 representation.

Two four-bit values fit inside one uint8.

The compiler still needs to know which logical dimension was compressed into that byte.

For a left-hand operand, lhs_k_pack=False means the two FP4 values are packed along M rather than K.

For a right-hand operand, rhs_k_pack=False means they are packed along N rather than K.

The default for both flags is true, meaning the packed pair extends K.

A physical tensor shape therefore does not completely describe the matrix.

Consider a right-hand FP4 tensor loaded with physical shape:

[32, 64] bytes

With:

rhs_k_pack=False

each byte expands along N, so its logical matrix shape is:

[32, 128] FP4 values

The physical K dimension remains 32.

With the default:

rhs_k_pack=True

the same physical shape would instead represent:

[64, 64] FP4 values

because the packed pair expands K.

The flag controls mathematical interpretation, not only memory performance.

AccelerateMatmul rewrites the dot through a transpose identity

The failing operation had:

A:
  BF16
  unscaled
  logical shape [128, 32]

B:
  E2M1 FP4
  scaled
  physical shape [32, 64]
  rhs_k_pack=False
  logical shape [32, 128]

output:
  [128, 128]

AccelerateMatmul identifies a dot_scaled whose left operand has no scale and whose right operand does.

It rewrites the multiplication by using:

(A × B)ᵀ = Bᵀ × Aᵀ

and then transposes the result back.

The transformed operation therefore has:

new lhs = Bᵀ
new rhs = Aᵀ
new lhs scale = old rhs scale
new rhs scale = old lhs scale
new lhs element type = old rhs element type
new rhs element type = old lhs element type

This moves the scaled FP4 operand from the right side to the left side while preserving the final matrix result.

The old implementation handled every item in that list except the packing flags.

A transpose changes which flag owns the packing contract

Before transposition, B’s two FP4 values were packed along N:

physical B: [32, 64]
logical B:  [32, 128]

After transposition, its physical shape becomes:

physical Bᵀ: [64, 32]

The former N dimension is now M.

The correct description is therefore:

lhs_k_pack=False

That tells Triton to expand the new left-hand operand along M:

physical Bᵀ: [64, 32]
logical Bᵀ:  [128, 32]

K remains 32, as required by the matrix identity.

The old transformation did not pass lhs_k_pack or rhs_k_pack to the new DotScaledOp.

Both attributes therefore used their default value:

true

For the new FP4 left operand, lhs_k_pack=True expands its final dimension:

physical Bᵀ: [64, 32]
misread as:  [64, 64]

The compiler had not moved or repacked a byte. It had changed the metadata that says what each byte means.

Why the scale factor became 64

The regression’s scale tensor has logical shape:

[N, K / group_size]
=
[128, 1]

The intended K dimension is 32, so the scale factor is:

32 / 1 = 32

Triton accepts scale factors of 16 or 32 for this operation.

Once the transposed FP4 input defaulted to K-packed form, Triton inferred:

logical K = physical K × 2
          = 32 × 2
          = 64

The same scale tensor still had one scale entry along K.

Validation therefore calculated:

64 / 1 = 64

and emitted the reported error.

The validator was not inventing the wrong number. It was consistently evaluating an incorrectly reconstructed operation.

The patch swaps rather than copies the flags

The corrected builder call adds:

/*lhs_k_pack=*/dotOp.getRhsKPack(),
/*rhs_k_pack=*/dotOp.getLhsKPack()

The sides must be exchanged because the matrix operands were exchanged.

No boolean inversion is required.

An original right-hand value of false means packing along N. After transposition, that dimension becomes M, and false on the new left-hand operand means packing along M.

The meaning moves with the data:

old RHS N-packed
        ↓ transpose and side swap
new LHS M-packed

Likewise, an old left-hand M-packed operand would become a new right-hand N-packed operand.

The same boolean value remains correct after it moves to the opposite attribute.

The structural test checks the lowering decision

The added MLIR test constructs the asymmetric case directly:

unscaled BF16 lhs
scaled E2M1 rhs
rhs_k_pack = false
CUDA target = sm_100

After TritonGPUAccelerateMatmul, the test requires the high-level tt.dot_scaled to be gone.

It also checks for:

ttg.fp4_to_fp ... axis = 0
ttng.tc_gen5_mma

The conversion axis is important. It shows that the transposed FP4 operand is being expanded along its non-K dimension rather than being interpreted as K-packed.

Reaching the tensor-core MMA operation also demonstrates that the pass no longer terminates at the invalid scale-factor check.

This test is focused on compiler structure. It verifies the relevant lowering path without depending on a full Python runtime invocation.

Review added an end-to-end numerical test

The first review requested an end-to-end case.

The final patch adds a Triton kernel that reproduces the same layout:

tl.dot_scaled(
    a,
    None,
    "bf16",
    b,
    b_scale,
    "e2m1",
    rhs_k_pack=False,
)

The test creates:

A: [128, 32] BF16
B: [32, 128] logical MXFP4, physically packed along N
scale: [128, 1]
output: [128, 128] FP32

It runs the compiled kernel and compares the result with a PyTorch reference that explicitly converts the FP4 values and scales to floating point before matrix multiplication.

The comparison uses:

atol = 1e-3
rtol = 1e-3

The test is enabled for NVIDIA CUDA devices with compute capability major version 8 or later.

After that test was added, the reviewer approved the change.

What the tests establish

The two tests cover different contracts.

The MLIR test establishes that:

the transpose rewrite retains the non-K packing choice
the scale factor remains valid
the operation reaches accelerated lowering

The Python test establishes that:

the compiled kernel executes
the packed FP4 operand is interpreted correctly
the numerical output matches an explicit reference

Together they protect both transformation structure and user-visible results.

This is stronger than testing only that compilation no longer crashes. A compiler could avoid the diagnostic while still unpacking the bytes along the wrong axis.

What remains outside the regression matrix

The committed regression is intentionally narrow.

It covers a right-hand e2m1 operand packed along N, with the left operand in BF16 and only the right operand scaled.

It does not directly test:

a mirrored left-hand non-K-packed input
both operands using explicit scales
batched dot shapes
AMD lowering
every supported microscaling format

The production change forwards both flags symmetrically, and the algebra of the swap supports that implementation. The strongest runtime evidence remains the reported asymmetric FP4 case.

The PR contains no performance benchmark, which is appropriate for a correctness fix whose observable failure was a compiler crash.

A small source comment remains confusing

The nearby helper comment says it transposes scaled dots that have a scale on the left-hand side.

The pass’s selection logic does the opposite: it collects operations whose A scale is absent and B scale is present, then transposes them so the scaled operand becomes the new left-hand side.

The code path and tests are consistent with the PR description. The comment appears stale.

That wording does not change the transformation, but correcting it would make the ownership of the scale and packing attributes easier to follow during future maintenance.

The compiler contract restored by the change

A transpose preserves a matrix expression only when every operand property is transformed with the operand.

That includes properties which do not appear in the physical tensor shape.

The corrected rewrite is:

transpose operands
        ↓
swap scales
        ↓
swap element formats
        ↓
swap packing-axis flags
        ↓
build the transformed dot
        ↓
transpose the result back

Copying values is not enough when metadata determines their logical dimensions.

Triton PR #11613 restores that contract for FP4 scaled dots: the packing description now follows the bytes through the transpose.