[GPU] Fix TritonEmitterTest.ScaledDotIsSupportedByReferencePlatform

[mmakevic-amd]

Motivation

This PR fixes TritonEmitterTest.ScaledDotIsSupportedByReferencePlatform subtest from //xla/backends/gpu/codegen/triton/tests:fusion_emitter_device_test.

Test started falling on ROCm after openxla@80931bb enforced Triton backend via:

debug_options.set_xla_gpu_experimental_disable_binary_libraries(true);

Until this point, the subtest was using hipblaslt instead of Triton. The only reason it was still passing on CUDA side was that CUBLASLT_FISSION was not added to the list of disabled backends. After changing that subtest consistently failed with mismatch errors on both platforms.

Failing log:

Mismatch count 256 (100.0000%) in shape bf16[16,16] (256 elements), abs bound 0.001, rel bound 0.001
Top relative error mismatches:
  actual       0.2617, expected   -6.407e-07, index {7,10}, rel error 4.08e+05, abs error    0.262
  actual       0.3223, expected   -2.205e-06, index {10,8}, rel error 1.46e+05, abs error    0.322
  actual       0.4082, expected    6.407e-06, index {2,11}, rel error 6.37e+04, abs error    0.408
  actual       0.2139, expected   -4.262e-06, index {8,12}, rel error 5.02e+04, abs error    0.214
  actual       0.3574, expected   -7.749e-06, index {8,6}, rel error 4.61e+04, abs error    0.357
Absolute magnitude breakdown of actual values:
  0      <= x < 0.0001 :       0 (  0.0000%)
  0.0001 <= x < 0.001  :       0 (  0.0000%)
  0.001  <= x < 0.01   :       0 (  0.0000%)
  0.01   <= x < 0.1    :       8 (  3.1250%), mismatches 8
  0.1    <= x < 1      :     248 ( 96.8750%), mismatches 248
  1      <= x < inf    :       0 (  0.0000%)
Elements exceeding abs error bound 0.001: 256 (100.0000%)
Relative error breakdown of elements exceeding abs error bound:
  <  0.0001 :       0 (0.0000%)
  >= 0.0001 :     256 (100.0000%)
  >= 0.001  :     256 (100.0000%)
  >= 0.01   :     256 (100.0000%)
  >= 0.1    :     256 (100.0000%)
  >= 1      :     256 (100.0000%)
Elements exceeding rel error bound 0.001: 256 (100.0000%)
Absolute error breakdown of elements exceeding rel error bound:
  <  0.0001 :       0 (0.0000%)
  >= 0.0001 :     256 (100.0000%)
  >= 0.001  :     256 (100.0000%)
  >= 0.01   :     256 (100.0000%)
  >= 0.1    :     248 (96.8750%)
  >= 1      :       0 (0.0000%)

Technical Details

XLA does not set scales for BF16 type https://github.com/openxla/xla/blob/111f921748b0643e70693b6b2a53a0d76b2af37a/xla/codegen/xtile/codegen/dot_algorithms.cc#L78-L120:

  Value lhs_scale;
  if (lhs_dot_elem_type != b.getBF16Type()) {
    lhs_scale = Bitcast(b, operands.lhs_scale, b.getI8Type());
  }
  Value rhs_scale;
  if (rhs_dot_elem_type != b.getBF16Type()) {
  ...

and it also skips ScaledDotRewriter for Triton https://github.com/openxla/xla/blob/111f921748b0643e70693b6b2a53a0d76b2af37a/xla/service/gpu/gpu_compiler.cc#L704-L706:

if (!debug_options.xla_gpu_experimental_scaled_dot_with_triton()) {
    pipeline.AddPass<ScaledDotRewriter>();
  }

So it emits a dot_scaled with no scales, which explains the accuracy issues. This PR modifies ScaledDot function so that for bf16 it converts scale to operand type, broadcasts + reshapes it to operand shape, multiplies, then emits a plain stablehlo.dot_general via EmitStableHloDotAndAdd. This follows the logic from ScaledDotRewriter::RewriteComputation (https://github.com/openxla/xla/blob/9e1e4f3a8d7db8c626cd4079058353ce87d4f941/xla/backends/gpu/transforms/scaled_dot_rewriter.cc#L166-L192)

Test Result

Problematic test now passes. After inspecting Triton dumps confirmed that instead of tt.dot_scaled, it now emits tt.dot.

Submission Checklist

• Look over the contributing guidelines at https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

[Copilot]

Pull request overview

This PR addresses correctness issues for GPU Triton lowering of scaled-dot when operands are BF16 (notably fixing TritonEmitterTest.ScaledDotIsSupportedByReferencePlatform) by avoiding emission of tt.dot_scaled without valid scales, and aligns autotuner backend disabling behavior when binary libraries are turned off.

Changes:

• Disable CUBLASLT_FISSION autotuning backend when xla_gpu_experimental_disable_binary_libraries is enabled.
• Add a BF16-specific scaled-dot lowering path that applies scales via broadcast/reshape + multiply and then emits a plain stablehlo.dot_general (+ add) instead of xtile::DotScaledOp.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
xla/service/gpu/autotuning/autotuner_pass.cc	Disables CUBLASLT_FISSION when binary libraries are disabled to prevent unintended backend selection.
xla/codegen/xtile/codegen/dot_algorithms.cc	Introduces BF16 scaled-dot handling by dequantizing via explicit mul with broadcast scales and emitting stablehlo.dot_general.

[claude]

Nit: The || means a mixed F8+BF16 scenario would take this dequantize path for both operands, even though only the BF16 side truly needs it (F8 is supported by tt.dot_scaled). In the reference ScaledDotRewriter, BF16 operands get dequantized individually while F8 operands can still go through the scaled-dot path. This is fine for the current BF16+BF16 test case, but worth keeping in mind if mixed-precision scaled dots become supported.

[mmakevic-amd]

Same as above; fixed

[claude]

LGTM. CUBLASLT_FISSION was clearly missing from this list alongside the existing CUBLASLT, CUDNN, HIPBLASLT, MIOPEN, and HIPBLASLT_FISSION entries.

[claude]

Review Summary

The approach is sound: BF16 operands can't go through tt.dot_scaled (which expects F8/sub-byte types), so dequantizing via scale * operand and falling back to a plain dot is the right fix. The CUBLASLT_FISSION addition to the disabled backends list is clearly correct.

Main concern: The new DequantizeOperand function doesn't handle scalar (rank-0) scales, which the reference ScaledDotRewriter::Dequantize does. This would crash if a scalar scale is ever passed. See inline comments for details and a secondary precision note about scale type conversion direction.

Overall this looks good for the BF16+BF16 case it targets. The inline comments flag edge cases worth hardening for future-proofing.

[i-chaochen]

Thanks! please address Claude review and upstream it with [ROCm] prefix. :)

[i-chaochen]

I guess your changes are upon this commit openxla@28f3032 ?

[mmakevic-amd]

I guess your changes are upon this commit openxla@28f3032 ?

Yes these changes are old and already present on this branch

[mmakevic-amd]

Thanks! please address Claude review and upstream it with [ROCm] prefix. :)

All done. I enabled some other subtests for ROCm - I will split that in a separate PR for upstream, but let me know if I should do the same here

[i-chaochen]

Thanks! please address Claude review and upstream it with [ROCm] prefix. :)

All done. I enabled some other subtests for ROCm - I will split that in a separate PR for upstream, but let me know if I should do the same here

Super! please put into separate PR for upstream.

[mmakevic-amd]

Hi @zoranjovanovic-ns could you please take a look at this before I open a PR to the upstream?