[Issue]: Multi-device setup reports device 0's name and architecture for all GPUs

[alexschroeter]

Problem Description

In a mixed-GPU system, jax.devices() reports the marketing name of device 0 for all devices. More critically, this appears to affect compilation — XLA compiles kernels for the wrong architecture on devices other than device 0.

Setup: gfx906 (AMD Radeon VII) as device 0, gfx908 (AMD Instinct MI100) as device 1.

Both devices visible:

>>> jax.devices()
[RocmDevice(id=0), RocmDevice(id=1)]
>>> [d.device_kind for d in jax.devices()]
['AMD Radeon VII', 'AMD Radeon VII']  # device 1 is actually MI100

Each device isolated:

# ROCR_VISIBLE_DEVICES=0
>>> jax.devices()[0].device_kind
'AMD Radeon VII'  # correct

# ROCR_VISIBLE_DEVICES=1
>>> jax.devices()[0].device_kind
'AMD Instinct MI100'  # correct

The naming issue isn't just cosmetic. When placing a matmul on device 1 (MI100) in multi-device mode, XLA fails with:

error: unsupported target: 'gfx906'
JaxRuntimeError: INTERNAL: Autotuning failed ... No valid config found!

This suggests XLA is compiling gfx906 kernels for the MI100. The same matmul works perfectly when the MI100 is isolated with ROCR_VISIBLE_DEVICES=1.

Non-GEMM operations (elementwise, reductions, FFT) do work on both devices in multi-device mode, though they emit Triton warnings like error: unsupported target: 'gfx906' — again suggesting wrong-arch compilation is being attempted.

Environment: ROCm 7.2.1 (rocm/dev-ubuntu-24.04:7.2.1-complete), JAX 0.9.1, jax-rocm7-plugin 0.9.1.post3, Ubuntu 22.04 host (kernel 6.8.0-87-generic).

Workaround: Use ROCR_VISIBLE_DEVICES to isolate a single GPU type per process.

Operating System

DOCKER:rocm/dev-ubuntu-24.04:7.2.1-complete

CPU

Intel(R) Xeon(R) Gold 6130 CPU @ 2.10GHz

GPU

MI100, Radeon VII

ROCm Version

7.2.1

ROCm Component

No response

Steps to Reproduce

No response

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

[lucbruni-amd]

Thanks for the report, @alexschroeter. I can reproduce many of these symptoms with a gfx906 (Radeon VII) as device 0 and a gfx1101 (Radeon PRO W7700) as device 1. Let me investigate what's going on and I'll update this issue when I have further findings.

[lucbruni-amd]

The linked PR in openxla/xla is to address the cosmetic issue.

The linked PR in rocm-systems, particularly clr, is to address the Matmul on device 1 fails with Autotuning failed ... No valid config found! error. From my investigation, this was not a wrong-target compilation bug in XLA, and the autotuner correctly selects the right device and architecture at every step. The actual bug is in HIP's ExtractFatBinaryUsingCOMGR which loads fat binary code objects for all devices at once. When it encounters a device with no compatible code object (gfx906 has no matching entry in the default build targets), it breaks out of the loop, preventing code objects from being loaded for remaining devices, even those with a valid match (gfx908 is an explicit build target). hipGetFuncBySymbol then returns error 98 for all devices, and the autotuner crashes because it can't initialize buffers.

Again ROCR_VISIBLE_DEVICES works because isolating a single GPU bypasses the multi-device loop entirely, allowing the matching code object to load successfully.