Problem Description
When a system has a GPU whose architecture is no longer supported by rocBLAS (e.g., gfx906 on ROCm 7.2+), JAX operations that dispatch to rocBLAS (matmul, GEMM) crash the entire process with SIGABRT. This is uncatchable from Python — there's no opportunity to fall back gracefully.
The crash happens at operation time, not at import. import jax, jax.devices(), array creation, and non-GEMM operations (elementwise, reductions, FFT) all work fine. The abort triggers when XLA's autotuner dispatches to rocBLAS, starting around 10x10 matmul:
rocBLAS error: Cannot read /opt/rocm/lib/rocblas/library/TensileLibrary.dat: Illegal seek for GPU arch : gfx906
Process dies immediately with exit code 134.
This is ultimately a rocBLAS issue — it calls abort() instead of returning an error code when Tensile kernel libraries are missing for the active GPU. For comparison, hipBLASLt already handles unsupported architectures gracefully by logging a warning and falling back. That said, it would be great if JAX/XLA could add a guard on the JAX side as well, since waiting for an upstream rocBLAS fix may take time and users have no way to protect themselves. If you feel I should open a ticket with rocBLAS since JAX is most likely not the only library suffering from this issue, please let me know.
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, gfx906 (Radeon VII) + gfx908 (MI100).
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
rocBLAS
Steps to Reproduce
No response
(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support
No response
Additional Information
No response
Problem Description
When a system has a GPU whose architecture is no longer supported by rocBLAS (e.g., gfx906 on ROCm 7.2+), JAX operations that dispatch to rocBLAS (matmul, GEMM) crash the entire process with SIGABRT. This is uncatchable from Python — there's no opportunity to fall back gracefully.
The crash happens at operation time, not at import. import jax, jax.devices(), array creation, and non-GEMM operations (elementwise, reductions, FFT) all work fine. The abort triggers when XLA's autotuner dispatches to rocBLAS, starting around 10x10 matmul:
rocBLAS error: Cannot read /opt/rocm/lib/rocblas/library/TensileLibrary.dat: Illegal seek for GPU arch : gfx906
Process dies immediately with exit code 134.
This is ultimately a rocBLAS issue — it calls abort() instead of returning an error code when Tensile kernel libraries are missing for the active GPU. For comparison, hipBLASLt already handles unsupported architectures gracefully by logging a warning and falling back. That said, it would be great if JAX/XLA could add a guard on the JAX side as well, since waiting for an upstream rocBLAS fix may take time and users have no way to protect themselves. If you feel I should open a ticket with rocBLAS since JAX is most likely not the only library suffering from this issue, please let me know.
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, gfx906 (Radeon VII) + gfx908 (MI100).
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
rocBLAS
Steps to Reproduce
No response
(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support
No response
Additional Information
No response