[Issue] JAX: `core_test` failures due to XLA autotuning `No valid config found` in portable wheel CI

[chiranjeevipattigidi]

Summary:
The “Release portable Linux JAX Wheels” CI job fails when running JAX tests on ROCm. Multiple tests in jax/tests/core_test.py fail with jax.errors.JaxRuntimeError: INTERNAL: Autotuning failed ... NOT_FOUND: No valid config found!.

Effected Archs: gfx1151

Failed Tests:

• jax/tests/core_test.py::CoreTest::test_simple_jit
• jax/tests/core_test.py::CoreTest::test_tracer_reprs
• jax/tests/core_test.py::CoreTest::test_verbose_tracer_reprs

Full Error:

• https://github.com/ROCm/TheRock/actions/runs/23786307643/job/69331875632

Error

Autotuning failed for HLO: %input_reduce_fusion = pred[] fusion(...), kind=kInput ... with error: NOT_FOUND: No valid config found!
Autotuning failed for HLO: %loop_reduce_fusion = f32[] fusion(...), kind=kLoop ... with error: NOT_FOUND: No valid config found!

[srinivamd]

Root Cause Summary

The CI failure is a two-part problem: a recent CI infrastructure fix exposed a pre-existing XLA architectural gap on gfx1151.

---

Root Cause

1. TheRock PR #4012 (merged 2026-03-27) added JAX_PLATFORMS=rocm to the JAX wheel test step, along with an /opt/rocm symlink for RUNPATH resolution. Before this fix, tests silently fell back to CPU and passed. After the fix, tests correctly run on GPU — and the real failure is exposed on gfx1151.

2. XLA's GPU reduction autotuner has no valid tile configurations for gfx1151 (STRIX-HALO, RDNA3.5). The underlying cause is in xla/stream_executor/rocm/rocm_compute_capability.h and xla/service/gpu/reduction_utils.cc:

• gfx1151 uses 32-thread wavefronts (RDNA) vs 64 threads on CDNA
• gfx9_mi100_or_later() returns false for gfx1151 — it only covers MI100/MI200/MI300 series
• has_mfma_instr_support() returns false for gfx1151
• reduction_utils.cc calls WarpSize(device_description) throughout its tile validity checks — this returns 32 for gfx1151, causing the tile constraint solver to exhaust its search space and return NOT_FOUND

The two failing fusion types in the error correspond to:

• kInput[pred[] input reduce fusion] — a boolean reduction (e.g. jnp.sum(x > 0))
• kLoop[f32[] loop reduce fusion] — a scalar float reduction (e.g. jnp.sum(x))

---

Prior Pattern

ROCm/jax PR #720 (merged 2026-03-04, "skip fused attention navi lds limit") already documented the same pattern of gfx1151/RDNA XLA feature gaps:

• Fused attention exceeds gfx1151's 64KB LDS limit (vs larger on CDNA) → test skipped
• TF32 dot unsupported (has_mfma_instr_support() = false on RDNA) → test skipped

The reduction autotuner failure is another instance of the same underlying issue: XLA code paths tuned for CDNA wavefront geometry (64 threads) don't produce valid configurations for RDNA's 32-thread wavefronts.

---

Recommended Actions

• Short-term (unblock CI): Add skip markers in ROCm/jax for the failing kInput/kLoop reduction fusion tests on gfx1151, following the same pattern as ROCm/jax#720. Target: skip on gfx11_apu() architectures (gfx1103, gfx1150, gfx1151).

• File a tracking issue in ROCm/jax and/or openxla/xla for gfx1151 reduction autotuner support, noting the 32-thread wavefront constraint.

• Long-term (proper fix): Add valid XLA reduction autotuner tile configurations for 32-thread RDNA wavefront geometry in openxla/xla. The tile constraint logic in xla/service/gpu/reduction_utils.cc needs to handle WarpSize == 32 and produce valid configs for gfx11 APUs.

• Diagnostic: Try XLA_FLAGS=--xla_gpu_autotune_level=0 (disable autotuning, use default configs) to confirm the issue is autotuner-specific and that tests pass with fallback configs.

[ckimamd]

@WBobby It looks like PR #4012 triggered these failures.
Could you please take a look?
CC: @amd-aakash @JeniferC99

[ckimamd]

core_test doesn't show up in the logs for Upstream JAX pytests on ROCm 7.2 using 1gpu, 4gpu, 8gpu across py3.11-py3.14

• list of failed tests from Upstream JAX pytests on ROCm 7.2 using 1gpu py3.14

core_test also doesn't show up in the fails, skipped, or error lists from runs on ROCm 7.12 or 7.2.1 on MI300X using ROCm/jax and jax-ml/jax v0.9.0:

• failed, skipped, errors for ROCm/jax on ROCm 7.12 MI300X
• failed, skipped, errors for jax-ml/jax on ROCm 7.12 MI300X
• failed, skipped, errors for ROCm/jax on ROCm 7.2.1 MI300X
• failed, skipped, errors for jax-ml/jax on ROCm 7.2.1 MI300X

[chiranjeevipattigidi]

linking other test failures in jax-gfx1151
#4373
#4374

[ckimamd]

Updated Investigation — 2026-04-07

Following up with a cross-repo status check. This issue is not yet resolved. Scope has expanded with two new sibling issues (#4373, #4374) filed today.

---

Resolution Status

Issue	Problem	Fix status
#4270 (this issue)	core_test.py reduction autotuner NOT_FOUND on gfx1151	Not fixed — no PR filed
#4373	scipy_stats_test.py device kernel image is invalid on gfx1151	Potential fix: ROCm/xla#780 (open, not merged)
#4374	multi_device_test.py abort — No ROCm wheel installation found	Not fixed — wheel discovery/packaging issue

---

Cross-Repo Findings

ROCm/jax: No PR has been filed to add skip markers for the failing core_test.py reduction tests on gfx1151, as recommended in the earlier analysis comment. The established skip pattern from ROCm/jax#720 (merged 2026-03-04, fused attention + TF32 skips for RDNA) is the right template — detection via "Radeon" in device_kind.

ROCm/xla: ROCm/xla#780 is a cherry-pick of openxla/xla#40419 into rocm-jaxlib-v0.9.2 that fixes the HSACO module cache using a freed pointer address as a key (causing stale/invalid kernel modules when memory is reused). This is the most likely cause of the device kernel image is invalid errors in #4373. It is currently open and blocked pending review.

jax-ml/jax / openxla/xla (upstream): No upstream fix or tracking issue has been filed for the 32-thread RDNA wavefront reduction autotuner gap. The tile constraint logic in xla/service/gpu/reduction_utils.cc needs to handle WarpSize == 32 to produce valid configs for gfx11 APUs, but this has not been reported upstream.

---

Recommended Next Steps

Short-term (CI unblock):

1. ROCm/jax: File a PR adding skip markers for core_test.py reduction fusion tests on gfx1151, following the ROCm/jax#720 pattern. Needs to land on both amd-main and rocm-jaxlib-v0.8.2 branches.
2. ROCm/xla: Merge ROCm/xla#780 to fix the device kernel image is invalid failures in [Issue] JAX ROCm wheel tests failing with invalid kernel file / device kernel image is invalid - gfx1151 #4373.
3. TheRock / ROCm/jax: Investigate wheel plugin discovery for multi_device_test.py ([Issue] JAX Wheels Test job aborts (exit 134) running jax/tests/multi_device_test.py due to missing ROCm wheel - gfx1151 #4374) — the jax_rocm7_pjrt plugin is not being found in the test environment.

Long-term (proper fix):
4. File a tracking issue in openxla/xla for gfx1151 reduction autotuner support — add valid tile configs for 32-thread RDNA wavefront geometry in xla/service/gpu/reduction_utils.cc.

[ckimamd]

Fix Status Update — 2026-04-07

Cross-repo search across ROCm/jax, jax-ml/jax, ROCm/xla, and openxla/xla confirms no single fix addresses all three linked issues. Each requires independent action:

#4270 — XLA reduction autotuner, no valid tile configs for gfx1151

• No PR filed in any repo. No skip markers filed in ROCm/jax. No upstream tracking issue filed in openxla/xla.
• Short-term: File skip markers in ROCm/jax for the failing core_test.py reduction tests on gfx1151, following the ROCm/jax#720 pattern.
• Long-term: File a tracking issue in openxla/xla for gfx1151 reduction autotuner support — add valid tile configs for 32-thread RDNA wavefront geometry in xla/service/gpu/reduction_utils.cc.

#4373 — HSACO module cache pointer bug (device kernel image is invalid)

• Upstream fix merged: openxla/xla#40419 / #40423 (2026-04-07)
• Short-term: Merge cherry-pick ROCm/xla#780 into rocm-jaxlib-v0.9.2 (open, pending review). File a follow-up cherry-pick to rocm-jaxlib-v0.8.2, which is the branch the TheRock wheel build defaults to.

#4374 — ROCm wheel/plugin not discoverable (No ROCm wheel installation found)

• No PR or fix filed in any repo.
• Short-term: Investigate wheel packaging and plugin path setup in the TheRock CI test step.