From eee3d454b6198e7e6bea2d5673891f6e7b10740a Mon Sep 17 00:00:00 2001 From: srinivamd <52507740+srinivamd@users.noreply.github.com> Date: Sun, 31 May 2026 21:24:10 -0700 Subject: [PATCH] [ROCm] Skip TF32 dot algorithm tests on Navi/RDNA GPUs TF32 (TensorFloat-32) is only supported on CDNA GPUs (MI100+/gfx9) via MFMA instructions. RDNA GPUs (Navi/Radeon) lack this hardware. Upstream PR jax-ml/jax#34534 changed the skip guard from unconditional to CUDA-only, which correctly enables TF32 tests on CDNA but also unblocks them on RDNA where XLA raises UNIMPLEMENTED. Cherry-pick the Navi/RDNA TF32 skip from rocm-jaxlib-v0.8.2 to skip TF32_TF32_F32 and TF32_TF32_F32_X3 on Radeon devices. Fixes: ROCM-25363 --- tests/lax_test.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/tests/lax_test.py b/tests/lax_test.py index f1aeec51ccdd..e73e1e582a19 100644 --- a/tests/lax_test.py +++ b/tests/lax_test.py @@ -1180,6 +1180,15 @@ def testDotAlgorithm(self, algorithm, dtype): raise SkipTest( f"The dot algorithm '{algorithm}' requires CUDA compute " "capability >= 8.0.") + # TF32 is not supported on Navi/RDNA GPUs. XLA only allows it on + # MI100+ (gfx9 CDNA) -- see xla/service/algorithm_util.cc. + if (algorithm in {lax.DotAlgorithmPreset.TF32_TF32_F32, + lax.DotAlgorithmPreset.TF32_TF32_F32_X3} + and jtu.is_device_rocm() + and "Radeon" in jax.local_devices()[0].device_kind): + raise SkipTest( + f"The dot algorithm '{algorithm}' is not supported on " + "Navi/RDNA GPUs.") elif algorithm not in { lax.DotAlgorithmPreset.DEFAULT, lax.DotAlgorithmPreset.ANY_F8_ANY_F8_F32, @@ -1189,15 +1198,6 @@ def testDotAlgorithm(self, algorithm, dtype): }: raise SkipTest( f"The dot algorithm '{algorithm}' is not supported on GPU.") - if jtu.test_device_matches(["tpu"]): - if algorithm not in { - lax.DotAlgorithmPreset.DEFAULT, - lax.DotAlgorithmPreset.BF16_BF16_F32, - lax.DotAlgorithmPreset.BF16_BF16_F32_X3, - lax.DotAlgorithmPreset.BF16_BF16_F32_X6, - }: - raise SkipTest( - f"The dot algorithm '{algorithm}' is not supported on TPU." ) lhs_shape = (3, 4) rhs_shape = (4, 3)