[ROCm] Insert ReshapeDecomposer before post-GemmRewriter LayoutNormalization

[magaonka-amd]

================== NOT MEANT FOR MERGING YET — DISCUSSION PR ==================

TL;DR ( if you dont want to read the whole story )

jax.nn.dot_product_attention(impl='xla') intermittently dies with a RET_CHECK at layout_normalization.cc:431 on ROCm under multi-GPU pytest. Root cause: when the autotuner picks hipBLASLt over Triton for a dot, GemmRewriter pins the __cublas$lt$matmul output to canonical {n-1,...,1,0}, but layout assignment had already planned the
downstream reshape for a non-canonical layout, the reshape is no longer a bitcast, and LayoutNormalization aborts. The flap is rare because hipBLASLt is normally slower than Triton, but contention inflates the autotuner's single-shot Triton timing while hipBLASLt's min(128 algos) stays stable.

The fix is two pipeline.AddPass<ReshapeDecomposer>() calls in gpu_compiler.cc.

---

Full details

Hey folks, I'm currently debugging an interesting JAX unit test issue where
tests/nn_test.py::NNFunctionsTest::testDotProductAttention* and
pallas/gpu_paged_attention_test.py::*test_quantized_paged_attention*
intermittently fail in pytest with:

    INTERNAL: RET_CHECK failure
        (external/xla/xla/service/layout_normalization.cc:431)
        ShapeUtil::ReshapeIsBitcast(s, operand->shape())

This is more likely to appear if we make pytest to see multiple GPUs. But it can be reproduced in single gpu run aswell. But it is not very easy to reproduce it needs certain condition to be met during autotune.

Below is full details:

The triggering JAX call is the standard jax.nn.dot_product_attention(impl='xla')
path. JAX traces a Grouped Query Attention computation which fissions into a
jit_dot_general HLO module containing a single batched dot_general. JAX side it seems alright to me.

Narrowing down to a minimal HLO

What I did is first tried to turn this flaky test into semi deterministic fail so that I can debug easily and I narrowed it down to below HLO

HloModule jit_dot_general, entry_computation_layout={(f16[2,1,32,4,128]{4,3,2,1,0}, f16[4,2,1,128,128]{4,3,2,1,0})->f16[2,1,32,128]{3,2,1,0}}

ENTRY %main.1 (args_0_.1: f16[2,1,32,4,128], args_1_.1: f16[4,2,1,128,128]) -> f16[2,1,32,128] {
  %args_0_.1 = f16[2,1,32,4,128]{4,3,2,1,0} parameter(0)
  %args_1_.1 = f16[4,2,1,128,128]{4,3,2,1,0} parameter(1)
  ROOT %dot_general.1 = f16[2,1,32,128]{3,2,1,0} dot(%args_0_.1, %args_1_.1),
       lhs_batch_dims={0,1}, lhs_contracting_dims={3,4},
       rhs_batch_dims={1,2}, rhs_contracting_dims={0,3}
}

just running this HLO once won't reproduce this issue you need to run this under parallel load only then in rare cases autotuner decides hipBLASLT is the winner based on profiling time and chooses it as backend. hipBLASLT has limitation.

Making the bug deterministic for debugging

I have hacked XLA to simulate same scenario and it can reproduce the bug reliably without needing to run it mutliple time.

you need to apply below patch :

diff --git a/xla/backends/gpu/autotuner/factory_rocm.cc b/xla/backends/gpu/autotuner/factory_rocm.cc
index 6eace69781..c9a70f3bcf 100644
--- a/xla/backends/gpu/autotuner/factory_rocm.cc
+++ b/xla/backends/gpu/autotuner/factory_rocm.cc
@@ -84,6 +84,23 @@ std::vector<std::unique_ptr<CodegenBackend>> GetCodegenBackendsForROCm(
     MLIRContext* mlir_context,
     absl::Span<const autotuner::Backend> backend_allowlist) {
   std::vector<std::unique_ptr<CodegenBackend>> backends;
+  {
+    bool enable_cublaslt_first = debug_options->xla_gpu_enable_cublaslt();
+    backends.push_back(std::make_unique<FissionBackend>(
+        debug_options, compiler, target_config,
+        std::make_unique<HipblasLtBackend>(stream_executor, debug_options,
+                                           compiler, target_config),
+        GetGemmRewriterPipeline(
+            target_config->device_description, /*enable_cublaslt=*/true,
+            enable_cublaslt_first ? absl::Span<const DType>(kAllDTypes)
+                                  : kFp8OnlyDTypes),
+        alias_info, mlir_context));
+  }
   backends.push_back(std::make_unique<TritonBackend>(
       debug_options, compiler, target_config, alias_info, mlir_context));
   backends.push_back(

Then run below command:

git apply patches/rocm-debug-force-cublaslt-first.patch
bazel build --config=rocm //xla/tools:run_hlo_module
HIP_VISIBLE_DEVICES=0 \
  XLA_FLAGS="--xla_gpu_autotune_level=0" \
  ./bazel-bin/xla/tools/run_hlo_module \
  --platform=ROCM --reference_platform="" \
  repro/ret_check_repro_hlo.txt

Why this layout mismatch is problematic

The catch with hipBLASLt is that it only knows how to write its output in one shape i,e canonical row-major.
When XLA's GemmRewriter rewrites a kDot into a __cublas$lt$matmul call, it has to pin the custom-call's
output layout to {n-1, n-2, ..., 1, 0} to match what the library expects. The compiler doesn't get a vote.

The problem is that layout assignment ran much earlier, and at that point it was free to choose any layout for the dot's output. It picked one that matched what the rest of the graph wanted: a downstream bitcast further along the chain expected to read its input in a non-canonical {2,3,1,0} layout, and layout assignment propagated that constraint backward through a reshape so that reshape would just be a free relabel.

Now hipBLASLt swaps in and forces the dot's output to {2,1,0}. The reshape suddenly has an operand laid out one way and an output expected the other way. It's no longer a free relabel to actually go from one
layout to the other you'd have to physically move bytes around. The next pass, LayoutNormalization, walks every kReshape and checks exactly this with ShapeUtil::ReshapeIsBitcast(...). It gets false back, realises the invariant is broken, and fires the RET_CHECK.

Concretely, here's the failing HLO snippet we captured (one pass before
the crash):

%custom-call.1 = (f32[8,128,32]{2,1,0}, s8[67108864]{0})
   custom-call(%reshape.13, %reshape.14),
   custom_call_target="__cublas$lt$matmul",
   backend_config={"selected_algorithm":"50",...}
%get-tuple-element = f32[8,128,32]{2,1,0} get-tuple-element(%custom-call.1), index=0
%reshape.15 = f32[2,4,128,32]{2,3,1,0} reshape(%get-tuple-element)   // issue
%bitcast.18 = f32[2,4,32,128]{3,2,1,0} bitcast(%reshape.15)

Why does the autotuner sometimes pick hipBLASLt over Triton?

Now we know the failure next part is to figure why we land in this scenario? why autotuner sometime picks hipBLASLT over Triton path:

For each dot, the autotuner asks every registered backend for its candidate list via GetSupportedConfigs(instr). For our shape, the Triton backend returns 1 candidate and the HipblasLt backend returns 128 candidates (one per algorithm enumerated by hipblasLtMatmulAlgoGetHeuristic, capped at the XLA constant GemmConfig::kNumAlgorithms=128). The autotuner then profiles all 129 candidates with 1 warmup + 1 timed run each and picks the one with minimum measured duration.

Generally this is how timing looks: we can clearly see Triton wins here

Backend	n	min (µs)	median (µs)	max (µs)	Notes
Triton	1	6.40	6.40	6.40	only 1 supported config
HipblasLt	128	24.84	31.80	156.88	128 algorithms

but under parallel workers sometime this flips like below: where hipBLASLT timing remain as it is but triton shoots up. Here comparison is unfair for triton where it gets single warmup and single timing while hipBLASLT gets 128 runs and we select min(128 runs time).
I'm taking extreme parallel case of 64 processes running the hlo to highlight the bug

Group	n	Triton min (µs)	Triton median	Triton max	HipblasLt min	HipblasLt median	HipblasLt max
failed	3	43.48	51.40	54.84	17.08	17.24	17.96
passed	61	6.08	6.40	13.64	15.80	17.40	20.08

Is the Triton kernel actually slower under load?

To measure this I fired up rocprofv3 and try to compare timing info in the autotuner

Contention	Triton GPU dispatch (rocprofv3)	Triton autotune-reported (XLA)
1 proc on GPU0	1.16 – 7.60 µs	6.40 µs
8 procs on GPU0	1.16 – 5.32 µs	10.08 µs
16 procs on GPU0	1.16 – 7.72 µs	8.60 µs

From the table it is clear that:
The actual Triton kernel runs in 1.2-7.7 µs every single time, regardless of contention. The "inflation" the autotuner sees is entirely in the wall-clock window between hipEventRecord(start) and hipEventRecord(stop) that XLA's profiler uses for timing. That window also catches whatever work the GPU did for OTHER processes between processing the start and stop
events, and hipEventElapsedTime reports it all as if it were our kernel runtime. The kernel is queued and waiting at the front of its stream while neighbor work gets time-sliced ahead of it. But I'm unable to prove this because Tracing itself adds enough serialization that the 50-100 µs spikes only show up without a profiler attached. I think it is some sort of Schrodinger cat scenario.

The fix

Inserts a ReshapeDecomposer pass before each post-GemmRewriter LayoutNormalization in gpu_compiler.cc. After GemmRewriter pins a __cublas$lt$matmul (hipBLASLt on ROCm) output layout, any downstream non-bitcast reshape is decomposed into transpose + bitcast, restoring the ShapeUtil::ReshapeIsBitcast invariant that LayoutNormalization
asserts and was crashing on under autotune-driven backend flips.

Some questions from me:

1. fix I'm proposing is reasonable? my rationale is even though hipBLASLT is slow it is still valid backend ( but with limitation) so XLA should be able to handle it and it should not fail with RET_CHECK(). very minimum let the autotuner take slower path and complete execution.

2. Should GpuLayoutAssignment be the real fix? A more thorough fix would be inxla/backends/gpu/transforms/layout_assignment.cc: when AddBackendConstraints processes a kDot that could be rewritten to __cublas$lt$matmul (any Triton-eligible dot, since the autotuner may later pick the FissionBackend(HipblasLt) candidate), pre-constrain the dot's output layout to {n-1,...,1,0} so the downstream consumer chain plans around it from the start. This is more invasive (it would affect every dot that could go cuBLAS-LT, not just the ones that actually do, potentially adding kCopy operations in the Triton-wins path).

3. Should the autotuner measurement methodology be improved?
   The root flakiness is that GpuProfiler::Profile does 1 warmup + 1 timed run per candidate using hipEventElapsedTime,
   which captures wait time from neighbor processes' work in multi-tenant GPU scenarios (the "exclusive lock" in
   GpuProfiler::Execute is intra-process only, GetGpuMutex is a static absl::Mutex per (Platform, device_ordinal), not an
   IPC lock). HipblasLt happens to be robust against this because its min(128 algo) is order-statistics-stable; but triton suffers it
   because of lack of potential candidate configs

[draganmladjenovic]

I was expecting you needed to make a change here https://github.com/openxla/xla/blob/main/xla/backends/gpu/autotuner/factory_rocm.cc#L58. Does test fail consistently with triton_gemm set to false https://github.com/openxla/xla/blob/main/xla/debug_options_flags.cc#L336

[magaonka-amd]

if you set triton_gemm set to false layout assignment takes into account make sure it works with hipBLASLT so no you wont see failuire.
it fails only when triton_gemm is set and cublasLT is set , then layout assignment writes generic output layout. IF autotuner picks hipBLASLT somehow, only then issue appears.

coming to fix in factory rocm, I didnt think of that before. I just tested it now:
it still gives me ret_check error but at different palce
RET_CHECK failure (thunk_emitter.cc:1378) LayoutUtil::LayoutsInShapesEqual(...)

I just added it here in factory_rocm.cc

std::unique_ptr<HloPassPipeline> GetGemmRewriterPipeline(
    const stream_executor::DeviceDescription& device_description,
    bool enable_cublaslt, absl::Span<const DType> dtypes) {
  auto pipeline = std::make_unique<HloPassPipeline>(
      enable_cublaslt ? "hipblaslt_rewriter_pipeline"
                      : "rocblas_rewriter_pipeline");
  pipeline->AddPass(std::make_unique<DotAlgorithmRewriter>());
  pipeline->AddPass(std::make_unique<ScaledDotRewriter>());
  for (DType dtype : dtypes) {
    GemmRewriterOptions options{dtype};
    options.enable_cublaslt = enable_cublaslt;
    auto gemm_rewriter = std::make_unique<GemmRewriter>(
        device_description.gpu_compute_capability(),
        device_description.runtime_version(), options);
    pipeline->AddPass(std::move(gemm_rewriter));
  }
  pipeline->AddPass(std::make_unique<ReshapeDecomposer>());  ----- HERE
  return pipeline;
}

[magaonka-amd]

it will fail consistently with the hacky patch I made , I added that in PR description , other than that it is always flaky.

[nurmukhametov]

One thing that concerns me here if I'm reading this right (please correct me if not). The hipBLASLt candidate is profiled as a __cublas$lt$matmul writing canonical {n-1,...,0}, and the layout-fixup transpose lives outside the profiled module, so the comparison is:

• Triton candidate: timed kernel produces the planned layout. Runtime cost = timed cost.
• hipBLASLt candidate: timed kernel produces canonical layout. Runtime cost = timed cost + a layout-fixup transpose nobody measured.

However, something similar could already happen elsewhere in autotuning where post-boundary passes add costs the profiler doesn't see.

[draganmladjenovic]

One thing that concerns me here if I'm reading this right (please correct me if not). The hipBLASLt candidate is profiled as a __cublas$lt$matmul writing canonical {n-1,...,0}, and the layout-fixup transpose lives outside the profiled module, so the comparison is:

* Triton candidate: timed kernel produces the planned layout. Runtime cost = timed cost.

* hipBLASLt candidate: timed kernel produces canonical layout. Runtime cost = timed cost + a layout-fixup transpose nobody measured.

However, something similar could already happen elsewhere in autotuning where post-boundary passes add costs the profiler doesn't see.

No I think it gets profiled as full gpu executable, so that is hipblast + anything else that triton might have fused but hipblaslt can't.

[claude]

Missing regression test. The PR doesn't include an HLO-level test to guard against this crash regressing.

A test using HloTestBase could construct the post-GemmRewriter HLO directly — with a __cublas$lt$matmul custom call already in place and the problematic reshape (canonical output layout {n-1,...,1,0} feeding a reshape that is no longer bitcast-compatible) — then run only the ReshapeDecomposer + LayoutNormalization portion of the pipeline. This avoids dependence on the autotuner selecting hipBLASLt and makes the test deterministic.

Even a simple test that verifies the module compiles without the RET_CHECK crash would be valuable.

[claude]

Nit: comment scope is slightly narrow. The comment mentions only __cublas$lt$matmul, but GemmRewriter also pins layouts for __cublas$gemm custom calls. Since ReshapeDecomposer is unconditional this doesn't affect correctness, but the comment could say "GEMM custom calls" to be more precise and avoid misleading a future reader into thinking this is cBLAS-LT-specific.

[claude]

Same comment scope nit applies here. Same suggestion as above — consider broadening "__cublas$lt$matmul" to "GEMM custom calls" in this second comment instance as well.

[claude]

Review Summary

The fix is mechanically correct and well-placed. Inserting ReshapeDecomposer before each LayoutNormalization call mirrors the existing pattern in RunLayoutNormalizationPasses() and correctly handles the case where GemmRewriter's layout pinning breaks the bitcast invariant.

Key observations:

• Both insertion sites are correct — they cover the two LayoutNormalization calls that follow GEMM rewriting in the post-layout-assignment pipeline.
• The fix is also valid on the CUDA path (cuBLAS LT pins layouts the same way), and the overhead of ReshapeDecomposer as a no-op is negligible.
• The PR is marked "NOT MEANT FOR MERGING YET" with open design questions about whether the root cause should instead be fixed in GpuLayoutAssignment and whether autotuner measurement methodology needs improvement for multi-tenant GPU scenarios.

Suggestions: A regression test and a minor comment scope improvement — see inline comments.

🤖 Generated with Claude Code