diff --git a/jax/_src/cudnn/scaled_matmul_stablehlo.py b/jax/_src/cudnn/scaled_matmul_stablehlo.py
index b6b80748fb24..ab064f492bfe 100644
--- a/jax/_src/cudnn/scaled_matmul_stablehlo.py
+++ b/jax/_src/cudnn/scaled_matmul_stablehlo.py
@@ -23,6 +23,7 @@
 from jax._src import core
 from jax._src import dispatch
 from jax._src import dtypes
+from jax._src import lax as lax_internal
 from jax._src import numpy as jnp
 from jax._src import tree_util
 from jax._src.custom_derivatives import custom_vjp
@@ -104,6 +105,23 @@ def _scaled_matmul_gpu_lowering(
   return [out.result]
 
 
+def _scaled_matmul_rocm_lowering(
+    ctx, a, b, a_scales, b_scales, preferred_element_type
+  ):
+  def _scaled_dot_lowering_impl(lhs, rhs, lhs_scales, rhs_scales):
+    return lax_internal.scaled_dot(
+        lhs,
+        rhs,
+        lhs_scale=lhs_scales,
+        rhs_scale=rhs_scales,
+        dimension_numbers=(((2,), (2,)), ((0,), (0,))),
+        preferred_element_type=preferred_element_type,
+    )
+  return mlir.lower_fun(_scaled_dot_lowering_impl, multiple_results=False)(
+      ctx, a, b, a_scales, b_scales
+  )
+
+
 def _scaled_matmul_abstract(a, b, a_scale, b_scale, *, preferred_element_type):
   batch, non_contracting_lhs, contracting_lhs = a.shape
   _, non_contracting_rhs, _ = b.shape
@@ -122,9 +140,14 @@ def _scaled_matmul_abstract(a, b, a_scale, b_scale, *, preferred_element_type):
     _scaled_matmul_gpu_lowering,
     platform="cuda",
 )
+
+# Keep CUDA lowering on the existing `__op$block_scaled_dot` custom call.
+# ROCm is registered separately because AMD's backend does not support that
+# custom call; instead, lowering through `lax.scaled_dot` lets XLA emit
+# `kScaledDot`, which ROCm can then fuse via Triton or hipBLASLt when possible.
 mlir.register_lowering(
     _scaled_matmul_p,
-    _scaled_matmul_gpu_lowering,
+    _scaled_matmul_rocm_lowering,
     platform="rocm",
 )
 
diff --git a/tests/ann_test.py b/tests/ann_test.py
index 18bb51bec93b..8114bbbda805 100644
--- a/tests/ann_test.py
+++ b/tests/ann_test.py
@@ -119,6 +119,10 @@ def test_autodiff(self, shape, dtype, k, is_max_k):
   )
   def test_pmap(self, qy_shape, db_shape, dtype, k, recall):
     num_devices = jax.device_count()
+    # TODO(araganes): Re-enable once upstream HloShardingV3 lands (JAX 0.9.2+).
+    # New pmap's SPMD tiling can't convert back to 1D pmap mesh on ROCm.
+    if jtu.is_device_rocm():
+      self.skipTest("IndivisibleError: SPMD tiling incompatible with 1D pmap mesh on ROCm")
     rng = jtu.rand_default(self.rng())
     qy = rng(qy_shape, dtype)
     db = rng(db_shape, dtype)