diff --git a/jax/_src/cudnn/scaled_matmul_stablehlo.py b/jax/_src/cudnn/scaled_matmul_stablehlo.py
index b6b80748fb24..b64a734bac74 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
@@ -124,7 +142,7 @@ def _scaled_matmul_abstract(a, b, a_scale, b_scale, *, preferred_element_type):
 )
 mlir.register_lowering(
     _scaled_matmul_p,
-    _scaled_matmul_gpu_lowering,
+    _scaled_matmul_rocm_lowering,
     platform="rocm",
 )
 
@@ -318,7 +336,6 @@ def _scaled_matmul_batcher(batched_args, batch_dims, *, preferred_element_type):
       and batch_dims[0] == batch_dims[2]
       and batch_dims[0] == batch_dims[3]
   )
-  lhs_bdims = batch_dims[0]
   out_bdims = (batch_dims[0],)
   lhs, rhs, lhs_scales, rhs_scales = batched_args
   *batch, lhs_non_contracting, contracting = lhs.shape
@@ -614,7 +631,6 @@ def scaled_dot_general_transpose_lhs(
   else:
     ans_batch, _, ans_y = ranges_like(x_batch, x_kept, y_kept)
 
-  dims = ((ans_y, y_kept), (ans_batch, y_batch))
   x_contract_sorted_by_y = list(np.take(x_contract, np.argsort(y_contract)))
   out_axes = np.argsort(list(x_batch) + x_kept + x_contract_sorted_by_y)
 
@@ -767,4 +783,4 @@ def scaled_dot_general_wrapper(
   # the output, which we need to squeeze.
   if dn_batched != dimension_numbers:
     return jnp.squeeze(out, axis=0)
-  return out
+  return out
\ No newline at end of file