OpenXLA PR #48664: Why BF16 Scales Disappeared Inside Triton dot_scaled
OpenXLA PR #48664 fixes a compiler path that approved a scaled matrix multiplication and then emitted an operation that could not carry part of its meaning.
The source HLO operation had four inputs:
left matrix
right matrix
left scale
right scaleFor BF16 operands, those scales could contain real values that had to be multiplied into the matrices before the dot product.
XLA’s Triton support check assumed that BF16 scales were always dummy values equal to one. It therefore allowed the operation without checking the scale contents.
The XTile emitter then created tt.dot_scaled, but attached scales only for the low-precision operand types that Triton actually dequantizes:
F4E2M1FN
F8E4M3FN
F8E5M2BF16 was not in that set. Its scale operand was omitted.
The compiler did not reject the program. It emitted an ordinary unscaled BF16 dot and silently produced the wrong result.
The patch changes the support decision: a scale that Triton will ignore is permitted only when XLA can prove that every value is one. Otherwise, the scaled dot is rejected from that backend and decomposed into explicit multiply operations plus a regular dot.
The source operation and target operation had different contracts
XLA’s scaled-dot is broader than Triton’s tt.dot_scaled.
At the HLO level, the scale tensors describe how each operand must be dequantized or scaled. The fallback rewriter can implement that meaning for ordinary floating-point inputs by:
convert operand and scale to a common type
↓
broadcast the scale to the operand shape
↓
multiply operand × scale
↓
execute an ordinary dotThe target Triton operation is more specialized.
XLA’s XTile emitter considers only three operand types to be scaled-dot types whose scales are attached to tt.dot_scaled:
F4E2M1FN
F8E4M3FN
F8E5M2A BF16 operand can participate in the operation, but it is treated as an already expanded, unscaled operand. No BF16 scale value is passed into the target operation.
That behavior is correct when the BF16 scale is only a placeholder of one. It is incorrect when the scale changes the operand.
The defect came from validating the source operation as though the target operation supported the same semantic surface.
The old support check trusted the operand type
Before the patch, IsTritonSupportedScaledDot() checked the scale type only when the associated operand was not BF16.
The source comment stated the assumption directly:
BF16 operands are unscaled
therefore their HLO scales are dummy constants such as 1.0Under that assumption, skipping the scale-type check looked harmless.
The check never proved the more important property:
Will discarding this scale leave the operand unchanged?A BF16 scale tensor could be a parameter, a broadcast constant, or another nontrivial value. The support function still returned “allowed.”
Later, the emitter used the actual operand type to decide whether to attach the scale. BF16 failed that test, so the scale disappeared.
Each local decision matched its own assumption:
support layer:
BF16 scale is probably a placeholder
emitter layer:
BF16 is not a Triton dequantized operand typeTogether they removed a real multiplication from the program.
Silent wrong code is more dangerous than a rejected lowering
A missing backend feature often produces a compile error.
This bug did not.
The resulting operation still had valid matrix shapes and supported BF16 types. Triton could compile and execute it. The only missing information was the scale that should have changed the numerical values.
The failure therefore looked like:
valid HLO scaled dot
↓
backend support returns allowed
↓
BF16 scales omitted during emission
↓
valid unscaled dot is generated
↓
execution completes with incorrect outputThere was no invalid pointer, unsupported instruction, or shape mismatch to force an obvious failure.
This is the compiler equivalent of dropping an argument that happens to be optional in the target API. The call still succeeds, but it no longer implements the source program.
The new rule asks whether the scale is observable
The corrected support logic separates two cases.
For operand types that Triton actually dequantizes, XLA verifies that the scale type is one of the supported scale formats.
For any other supported operand type, including BF16, XLA asks whether the scale is provably all ones.
Conceptually:
if target operation consumes this scale:
validate the scale representation
else:
require scale == 1 everywhereIf the scale is not provably one, the support decision returns a reason such as:
LHS scale is ignored for operand type BF16 but is not all ones.The operation is no longer eligible for direct Triton scaled-dot emission.
This is a semantic support check rather than a type-only support check. It tests whether the target operation can preserve the source operation’s result.
IsAllOnesScale follows the value back to its source
A dummy scale is not always represented as a constant directly connected to the dot.
Compiler transformations may wrap it in operations such as:
broadcast
reshape
bitcast
convert
copyA scale inside a fusion can also appear as a parameter whose actual operand is defined outside that fusion.
The new IsAllOnesScale() helper walks backward through those value-preserving operations. When it reaches a fusion parameter, it follows the parameter to the corresponding fusion operand. The proof succeeds only when the chain eventually reaches a constant literal whose elements are all one.
The analysis is intentionally conservative.
A dynamically computed tensor that is mathematically equal to one but cannot be proven from this limited chain is rejected from the direct scaled-dot path. XLA can still preserve correctness by using the decomposition fallback.
That tradeoff is appropriate here. Failing to prove that an ignored input is neutral should cost an optimization opportunity, not numerical correctness.
Unsupported scaled dots are decomposed instead of discarded
The GPU compiler already runs ScaledDotRewriter for scaled dots that the selected backend does not support.
The rewriter expands each non-scalar scale into the operand’s shape, multiplies it into the operand, converts both sides to a compatible type, and creates a regular dot.
After this patch, a BF16 scaled dot with real scale values follows that path:
BF16 lhs × broadcast(lhs_scale)
↓
BF16 rhs × broadcast(rhs_scale)
↓
regular dotThe scale is not forced into a target operation that ignores it. Its effect is materialized before the dot.
A BF16 scaled dot whose scale is provably all ones can still use the compact Triton path because omitting multiplication by one preserves the result.
The patch therefore does not disable BF16 participation in Triton scaled dots. It distinguishes a neutral placeholder from an observable scale.
The emitter adds a second guard
The final patch also checks omitted scales inside the XTile emitter.
That layer normally sees a fusion that has already passed backend support analysis. During fusion extraction, however, a constant defined outside the fusion can become a parameter in the extracted module. The local emitter may no longer be able to trace that parameter back to the original constant.
The emitter therefore accepts an omitted scale when either:
IsAllOnesScale(scale)or:
the scale has an effectively scalar shapeThe scalar-shape exception reflects the pipeline’s placeholder convention after extraction. Its safety depends on the earlier support decision having admitted only neutral omitted scales in the original module.
For a non-scalar scale that is neither consumed nor provably one, the emitter returns an invalid-argument error rather than silently continuing.
This is defense in depth. The primary policy lives in backend support; the emitter also refuses to manufacture wrong code when its local representation still exposes the mismatch.
Why ROCm exposed the problem first
The imported change is labeled as a ROCm fix because the execution regression was visible there.
The pull request explains that NVIDIA commonly avoided the faulty Triton path through a cuBLASLt fission fallback, while the corresponding hipBLASLt fission route did not survive the tested configuration. When that alternative backend was absent, ROCm reached the Triton lowering that dropped the scales.
The underlying defect was not AMD-specific.
The support check and XTile emission logic were shared. Any configuration that selected the same Triton path for BF16 operands with nontrivial scales could produce the same unscaled dot.
ROCm changed which path was reachable. It did not change the meaning of the missing multiplication.
The execution test now runs on ROCm
The changed device test uses:
lhs: BF16[16, 128]
rhs: BF16[128, 16]
lhs_scale: BF16[1, 4]
rhs_scale: BF16[4, 1]Those scale tensors are parameters, not constants known to contain one. Their shapes are also not scalar placeholders.
Before this patch, the test was skipped on ROCm with the statement that scaled dot was unsupported on AMD.
The skip is removed. The full pipeline now compares the generated result against the reference implementation on ROCm as well.
For this input, the corrected support decision cannot prove the BF16 scales are all ones, so the direct Triton scaled-dot form is forbidden and the rewriter must preserve them explicitly.
The test does not merely check that the support function returns a different string. It checks the numerical result after compilation and execution on the previously skipped platform.
An existing support test covers the decision boundary
The patch does not add a new test file for IsAllOnesScale().
OpenXLA already has a parameterized ScaledDotOperandTypes support test. It constructs operand and scale parameters and exercises the supported operand-type matrix.
For low-precision operand types that tt.dot_scaled dequantizes, compatible scale types remain eligible.
For BF16 with an unknown non-scalar scale parameter, the new policy must reject direct emission because the value cannot be proven neutral.
This existing test covers the changed backend decision. The re-enabled ROCm device test covers the resulting execution path.
The helper’s individual traversal cases—such as nested broadcasts, converts, and fusion-parameter chains—are not each introduced as separate unit tests in this pull request.
Review changed the name, not the policy
The public review discussion was brief.
A reviewer requested a more specific name for the helper rather than a generic predicate name. The author renamed it to:
IsAllOnesScaleThe reviewer then approved the change.
The naming matters because the helper is not a general constant evaluator. It proves one narrow property through a limited set of operations for a scale-specific correctness gate.
No public benchmark was provided, and none is required to establish this correctness repair.
What the patch establishes
The merged code supports the following conclusion:
XLA scaled-dot may carry real scales for BF16 operands
↓
Triton tt.dot_scaled does not consume those BF16 scales
↓
therefore direct emission is correct only when each omitted scale is neutral
↓
non-neutral or unproven scales must use explicit multiplication and a regular dotThe patch preserves the optimized path for provable all-one placeholders while preventing nontrivial scales from disappearing.
It does not make Triton natively support BF16 block scales. It prevents XLA from claiming that support when the generated operation cannot express it.
The compiler contract restored by the change
A backend support decision must describe the operation that will actually be emitted.
The corrected sequence is:
inspect the source operand and scale
↓
ask whether the target operation consumes that scale
↓
if not, prove that omitting it is neutral
↓
otherwise select a semantics-preserving fallback
↓
emit only after the source and target contracts agreeAn input may be omitted only when the compiler can prove that it has no effect.
OpenXLA PR #48664 restores that rule for BF16 scaled dots and prevents a valid compiler path from silently changing the program’s arithmetic.
Sources
https://github.com/openxla/xla/pull/48664
https://github.com/openxla/xla/pull/48543
https://github.com/openxla/xla/commit/22f6d7f22819c6956354abbaa09b9faf17ab02c6