Repository: jax-ml/jax
Pull request #25096 (2024-11-25 19:46:12)

Update ROCm scripts to match new build.py usage
This PR updates the ROCm scripts that use `build/build.py` to match its new usage that was added in https://github.com/jax-ml/jax/pull/24695


Patch:
diff --git a/build/rocm/dev_build_rocm.py b/build/rocm/dev_build_rocm.py
index 2be64152f..aa5754b78 100755
--- a/build/rocm/dev_build_rocm.py
+++ b/build/rocm/dev_build_rocm.py
@@ -77,13 +77,14 @@ def build_jax_xla(xla_path, rocm_version, rocm_target, use_clang, clang_path):
     build_command = [
         "python3",
         "./build/build.py",
-        "--enable_rocm",
-        "--build_gpu_plugin",
-        "--gpu_plugin_rocm_version=60",
+        "build"
         f"--use_clang={str(use_clang).lower()}",
+        "--wheels=jaxlib,jax-rocm-plugin,jax-rocm-pjrt"
+        "--rocm_path=%/opt/rocm-{rocm_version}/",
+        "--rocm_version=60",
         f"--rocm_amdgpu_targets={rocm_target}",
-        f"--rocm_path=/opt/rocm-{rocm_version}/",
         bazel_options,
+        "--verbose"
     ]
 
     if clang_option:
diff --git a/build/rocm/tools/build_wheels.py b/build/rocm/tools/build_wheels.py
index deb6ab703..ec825f40b 100644
--- a/build/rocm/tools/build_wheels.py
+++ b/build/rocm/tools/build_wheels.py
@@ -93,11 +93,12 @@ def build_jaxlib_wheel(
     cmd = [
         "python",
         "build/build.py",
-        "--enable_rocm",
-        "--build_gpu_plugin",
-        "--gpu_plugin_rocm_version=60",
+        "build"
+        "--wheels=jaxlib,jax-rocm-plugin,jax-rocm-pjrt"
         "--rocm_path=%s" % rocm_path,
+        "--rocm_version=60",
         "--use_clang=%s" % use_clang,
+        "--verbose"
     ]
 
     # Add clang path if clang is used.
diff --git a/jax/_src/api.py b/jax/_src/api.py
index 308e7c230..8d464b51d 100644
--- a/jax/_src/api.py
+++ b/jax/_src/api.py
@@ -2777,6 +2777,7 @@ def clear_backends():
   dispatch.xla_primitive_callable.cache_clear()
   util.clear_all_caches()
   pjit._infer_params_cached.cache_clear()
+  pjit._pjit_lower_cached.cache_clear()
   pjit._create_pjit_jaxpr.cache_clear()  # pytype: disable=attribute-error
   pjit._cpp_pjit_cache_fun_only.clear()
   pjit._cpp_pjit_cache_explicit_attributes.clear()
diff --git a/jax/_src/cudnn/fused_attention_stablehlo.py b/jax/_src/cudnn/fused_attention_stablehlo.py
index ef4e33ad0..8ccf08ec6 100644
--- a/jax/_src/cudnn/fused_attention_stablehlo.py
+++ b/jax/_src/cudnn/fused_attention_stablehlo.py
@@ -1022,9 +1022,11 @@ def dot_product_attention(query: Array,
   Returns:
     Output of the same shape as the query.
   """
-  # TODO(b/380898464): Check the compute capability, e.g., require GPU device,
-  # in the kernel implementation (c++) code.
+  # check if cuDNN is installed
   cudnn_version = check_cudnn_version()
+  # only support at least Ampere
+  if not check_compute_capability("8.0"):
+    raise RuntimeError("Require at least Ampere arch to run")
   layout = _normalize_layout(qkv_layout)
   if has_padding(mask_type) and (q_seqlen is None or kv_seqlen is None):
     raise ValueError("Require q_seqlen and kv_seqlen to generate padding mask")
@@ -1045,7 +1047,7 @@ def dot_product_attention(query: Array,
 
   # combine bias and mask
   if bias is None:
-    bias = mask
+      bias = mask
   else:
     if mask is not None:
       # should be broadcast to same shape
diff --git a/jax/_src/pjit.py b/jax/_src/pjit.py
index fb60f9d52..e0af0b7d2 100644
--- a/jax/_src/pjit.py
+++ b/jax/_src/pjit.py
@@ -1665,8 +1665,7 @@ def _pjit_call_impl_python(
   compiler_options_kvs = compiler_options_kvs + tuple(pgle_compile_options.items())
   # TODO(patrios): Do not pass mutable profile session through cached lowering
   # chain. Instead we need to move profilers dictionary to pxla module and use
-  # module as key. Right now we can't do that since there is no way to evict
-  # _pjit_lower_cached cache for in PGLE mode.
+  # module as key. Right now we can't do that since there is no way to evict _pjit_lower_cached cache for in PGLE mode.
   compiled = _resolve_and_lower(
       args, jaxpr=jaxpr, in_shardings=in_shardings,
       out_shardings=out_shardings, in_layouts=in_layouts,
@@ -1777,7 +1776,12 @@ def _pjit_call_impl(*args, jaxpr,
 pjit_p.def_impl(_pjit_call_impl)
 
 
-def _pjit_lower(
+def _pjit_lower(*args, **kwargs):
+  return _pjit_lower_cached(*args, **kwargs)
+
+
+@weakref_lru_cache
+def _pjit_lower_cached(
     jaxpr: core.ClosedJaxpr,
     in_shardings,
     out_shardings,
diff --git a/jax/experimental/pallas/ops/gpu/attention_mgpu.py b/jax/experimental/pallas/ops/gpu/attention_mgpu.py
index c4ac7e625..1c5b4d9f7 100644
--- a/jax/experimental/pallas/ops/gpu/attention_mgpu.py
+++ b/jax/experimental/pallas/ops/gpu/attention_mgpu.py
@@ -59,27 +59,17 @@ def attention(q, k, v, config: TuningConfig):
   q_heads_per_kv_head = num_q_heads // num_kv_heads
   if head_dim % 64:
     raise ValueError(f"{head_dim=} must be divisible by 64")
+  if batch_size != 1:
+    raise NotImplementedError(f"Only batch_size=1 is supported, got: {batch_size=}")
   if jnp.dtype(dtype) not in map(jnp.dtype, [jnp.float16, jnp.bfloat16]):
     raise NotImplementedError(f"Only f16 and bf16 are supported, got dtype: {dtype}")
-  # Squash batch and sequence dimensions.
-  # This is required because CUDA grid/TMA descriptors have a limited number of
-  # slice dimensions.
-  # TODO(apaszke): Implement slice squashing for TMAs.
-  q = jnp.reshape(q, (batch_size * q_seq_len, num_q_heads, head_dim))
-  k = jnp.reshape(k, (batch_size * kv_seq_len, num_kv_heads, head_dim))
-  v = jnp.reshape(v, (batch_size * kv_seq_len, num_kv_heads, head_dim))
-
+  q, k, v = map(lambda x: x[0], (q, k, v))
   max_concurrent_steps = min(
       config.max_concurrent_steps, kv_seq_len // config.block_kv
   )
   block_q, block_kv = config.block_q, config.block_kv
-  num_q_tiles, rem = divmod(q_seq_len, block_q * 2)
-  if rem:
-    raise NotImplementedError(f"{q_seq_len=} must be a multiple of {block_q * 2=}")
 
   def kernel(q_ref, k_ref, v_ref, out_ref, scoped):
-    bidx = lax.div(lax.axis_index("bq"), num_q_tiles)
-    qidx = lax.rem(lax.axis_index("bq"), num_q_tiles)
     smem_buffers, buffer_barriers, consumed_barriers, schedule_barrier = scoped
     wg_idx = lax.axis_index("wg")
     qo_smem2, k_smem, v_smem = smem_buffers
@@ -93,7 +83,7 @@ def attention(q, k, v, config: TuningConfig):
     def _compute_wg():
       plgpu.set_max_registers(232, action="increase")
       qo_smem = qo_smem2.at[wg_idx]
-      q_seq_base = qidx * (2 * block_q) + wg_idx * block_q + bidx * q_seq_len
+      q_seq_base = lax.axis_index("q") * (2 * block_q) + wg_idx * block_q
       q_head = lax.axis_index("heads")
 
       plgpu.copy_gmem_to_smem(
@@ -175,16 +165,14 @@ def attention(q, k, v, config: TuningConfig):
       plgpu.set_max_registers(40, action="decrease")
       kv_head = lax.div(lax.axis_index("heads"), q_heads_per_kv_head)
       for i in range(max_concurrent_steps):
-        start = i * block_kv + bidx * kv_seq_len
-        s = (pl.ds(start, block_kv), kv_head)
+        s = (pl.ds(i * block_kv, block_kv), kv_head)
         plgpu.copy_gmem_to_smem(k_ref.at[s], k_smem.at[i], k_barriers.at[i])
         plgpu.copy_gmem_to_smem(v_ref.at[s], v_smem.at[i], v_barriers.at[i])
 
       def kv_loop(kv_step, _):
         tma_step = kv_step + max_concurrent_steps
         tma_slot = lax.rem(kv_step, max_concurrent_steps)
-        start = tma_step * block_kv + bidx * kv_seq_len
-        s = (pl.ds(start, block_kv), kv_head)
+        s = (pl.ds(tma_step * block_kv, block_kv), kv_head)
         plgpu.barrier_wait(k_consumed_barrier)
         plgpu.copy_gmem_to_smem(k_ref.at[s], k_smem.at[tma_slot], k_barriers.at[tma_slot])
         plgpu.barrier_wait(v_consumed_barrier)
@@ -199,10 +187,13 @@ def attention(q, k, v, config: TuningConfig):
   def run(refs):
     q_ref, k_ref, v_ref, out_ref = refs
 
+    num_q_tiles, rem = divmod(q_seq_len, block_q * 2)
+    if rem:
+      raise NotImplementedError(f"{q_seq_len=} must be a multiple of {block_q * 2=}")
     mesh = plgpu.GPUMesh(
-        grid=(batch_size * num_q_tiles, num_q_heads),
+        grid=(num_q_tiles, num_q_heads),
         num_threads=3,
-        axis_names=("bq", "heads", "wg"),
+        axis_names=("q", "heads", "wg"),
         approx_math=True,
     )
     @pl.core_map(mesh)
@@ -236,7 +227,7 @@ def attention(q, k, v, config: TuningConfig):
       )
 
   _, _, _, out = pl.run_state(run)((q, k, v, jnp.full_like(q, jnp.inf)))
-  return jnp.reshape(out, [batch_size, q_seq_len, num_q_heads, head_dim])
+  return out[None]
 
 
 @jax.jit
@@ -256,13 +247,13 @@ def attention_reference(q, k, v):
 
 
 def main(unused_argv):
+  batch_size = 1
   num_q_heads = 1
   num_kv_heads = 1
-  problem_it = itertools.product((1, 2), (4096, 32768,), (64, 128, 256,))
-  for batch_size, seq_len, head_dim in problem_it:
+  problem_it = itertools.product((4096, 32768,), (64, 128, 256,))
+  for seq_len, head_dim in problem_it:
     q_seq_len = kv_seq_len = seq_len
-    print(f"==== {batch_size=:<6} {kv_seq_len=:<6} {q_seq_len=:<6}"
-          f"{num_q_heads=:<4} {head_dim=:<6} ====")
+    print(f"==== {kv_seq_len=:<6} {q_seq_len=:<6} {num_q_heads=:<4} {head_dim=:<6} ====")
     param_it = itertools.product((64,), (64, 128, 256))
     best = None
     k1, k2, k3 = jax.random.split(jax.random.key(42), 3)
diff --git a/jax/experimental/pjit.py b/jax/experimental/pjit.py
index 8ba7eb25d..a711c6bc4 100644
--- a/jax/experimental/pjit.py
+++ b/jax/experimental/pjit.py
@@ -22,3 +22,5 @@ from jax._src.sharding_impls import (
   AUTO as AUTO,
   UNSPECIFIED as _UNSPECIFIED,
 )
+
+from jax._src.pjit import _pjit_lower_cached, _pjit_lower


Test patch:
diff --git a/tests/BUILD b/tests/BUILD
index f0668b42b..a645a971a 100644
--- a/tests/BUILD
+++ b/tests/BUILD
@@ -260,6 +260,7 @@ jax_multiplatform_test(
     ],
 )
 
+# TODO(b/355263220): enable on TPU once layouts is supported with Shardy.
 jax_multiplatform_test(
     name = "layout_test",
     srcs = ["layout_test.py"],
@@ -278,6 +279,9 @@ jax_multiplatform_test(
 jax_multiplatform_test(
     name = "shard_alike_test",
     srcs = ["shard_alike_test.py"],
+    disable_configs = [
+        "cpu_shardy",  # TODO(b/355263220): enable once shard_alike is supported.
+    ],
     enable_configs = [
         "tpu_v3_2x2",
         "tpu_v5e_4x2",
diff --git a/tests/pjit_test.py b/tests/pjit_test.py
index e5aa1a604..4d9b98b4d 100644
--- a/tests/pjit_test.py
+++ b/tests/pjit_test.py
@@ -52,6 +52,7 @@ from jax._src import sharding_impls
 from jax._src.sharding_impls import (
     AUTO, UNSPECIFIED, NamedSharding, GSPMDSharding, PositionalSharding,
     SingleDeviceSharding, parse_flatten_op_sharding)
+import jax._src.pjit as pjit_lib
 from jax._src.pjit import pjit
 from jax._src import mesh as mesh_lib
 from jax._src.interpreters import pxla
@@ -1899,6 +1900,11 @@ class ArrayPjitTest(jtu.JaxTestCase):
   )
   def test_pjit_array_multi_input_multi_output(self, mesh_shape, s1_shape,
                                                s2_shape, s3_shape, s4_shape):
+    if config.use_shardy_partitioner.value:
+      self.skipTest(
+          'TODO(b/355263220) Shardy conflict resolution is not complete. Issue '
+          'here is that for `a1 @ a1.T` GSPMD gives dim 0 sharded on `x` while '
+          'Shardy gives it fully replicated.')
     global_mesh = jtu.create_mesh(mesh_shape, ('x', 'y'))
     global_input_shape = (8, 2)
 
@@ -2156,13 +2162,13 @@ class ArrayPjitTest(jtu.JaxTestCase):
       return x + y
 
     out = add(a, b)
-    cache_info1 = pxla._cached_lowering_to_hlo.cache_info()
+    cache_info1 = pjit_lib._pjit_lower_cached.cache_info()
     self.assertIsInstance(out, array.ArrayImpl)
     self.assertArraysEqual(out, a + b)
     self.assertFalse(out._committed)
 
     out2 = add(out, out)
-    cache_info2 = pxla._cached_lowering_to_hlo.cache_info()
+    cache_info2 = pjit_lib._pjit_lower_cached.cache_info()
     self.assertIsInstance(out2, array.ArrayImpl)
     self.assertArraysEqual(out2, 2 * (a + b))
     self.assertFalse(out2._committed)
@@ -2172,7 +2178,7 @@ class ArrayPjitTest(jtu.JaxTestCase):
 
     c = jax.device_put(a, jax.devices()[0])
     out3 = add(c, c)
-    cache_info3 = pxla._cached_lowering_to_hlo.cache_info()
+    cache_info3 = pjit_lib._pjit_lower_cached.cache_info()
     self.assertArraysEqual(out3, 2 * c)
     self.assertTrue(out3._committed)
 
@@ -2215,11 +2221,14 @@ class ArrayPjitTest(jtu.JaxTestCase):
 
     f = pjit(lambda x: x)
 
-    with jtu.count_jit_compilation_cache_miss() as count:
-      out1 = f(a)
-      out2 = f(b)
-    self.assertEqual(count[0], 2)
+    out1 = f(a)
+    cache_info1 = pjit_lib._pjit_lower_cached.cache_info()
 
+    out2 = f(b)
+    cache_info2 = pjit_lib._pjit_lower_cached.cache_info()
+
+    self.assertEqual(cache_info2.hits, cache_info1.hits)
+    self.assertEqual(cache_info2.misses, cache_info1.misses + 1)
     self.assertArraysEqual(out1, val1)
     self.assertArraysEqual(out2, val2)
 
@@ -2876,13 +2885,13 @@ class ArrayPjitTest(jtu.JaxTestCase):
       return x, y, z
 
     o1, o2, o3 = f(a, y=b, z=c)
-    cache_info1 = pxla._cached_lowering_to_hlo.cache_info()
+    cache_info1 = pjit_lib._pjit_lower_cached.cache_info()
     self.assertArraysEqual(o1, a)
     self.assertArraysEqual(o2, b)
     self.assertArraysEqual(o3, c)
 
     o4, o5, o6 = f(x=a, y=b, z=c)
-    cache_info2 = pxla._cached_lowering_to_hlo.cache_info()
+    cache_info2 = pjit_lib._pjit_lower_cached.cache_info()
     self.assertArraysEqual(o4, a)
     self.assertArraysEqual(o5, b)
     self.assertArraysEqual(o6, c)
@@ -2891,7 +2900,7 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertEqual(cache_info2.misses, cache_info1.misses + 1)
 
     o7, o8, o9 = f(a, b, c)
-    cache_info3 = pxla._cached_lowering_to_hlo.cache_info()
+    cache_info3 = pjit_lib._pjit_lower_cached.cache_info()
     self.assertArraysEqual(o7, a)
     self.assertArraysEqual(o8, b)
     self.assertArraysEqual(o9, c)
@@ -2978,19 +2987,26 @@ class ArrayPjitTest(jtu.JaxTestCase):
     x = jnp.arange(8).reshape(4, 2)
     f_out = f(x)
     f_out2 = f(f_out)
+    cache_info1 = pjit_lib._pjit_lower_cached.cache_info()
     _check(f_out, jax.devices()[1], x)
     _check(f_out2, jax.devices()[1], f_out)
 
     y = jax.device_put(x, jax.sharding.NamedSharding(mesh, P('x', 'y')))
     out2 = f(y)
+    cache_info2 = pjit_lib._pjit_lower_cached.cache_info()
     _check(out2, jax.devices()[1], y)
 
+    self.assertEqual(cache_info2.hits, cache_info1.hits + 1)
+
     with jtu.ignore_warning(category=DeprecationWarning,
                             message="backend and device argument"):
       h = pjit(mul, device=jax.devices()[-1])
     h_out = h(y)
+    cache_info3 = pjit_lib._pjit_lower_cached.cache_info()
     _check(h_out, jax.devices()[-1], y)
 
+    self.assertEqual(cache_info3.hits, cache_info2.hits)
+
     # AOT test
     compiled = f.lower(core.ShapedArray(y.shape, y.dtype)).compile()
     out3 = compiled(y)
@@ -3520,11 +3536,11 @@ class ArrayPjitTest(jtu.JaxTestCase):
 
     with jtu.count_pjit_cpp_cache_miss() as count:
       out = f(arr)
-      cache_info1 = pxla._cached_lowering_to_hlo.cache_info()
+      cache_info1 = pjit_lib._pjit_lower_cached.cache_info()
       self.assertIsInstance(out.sharding, NamedSharding)
 
       out2 = f(np_arr)
-      cache_info2 = pxla._cached_lowering_to_hlo.cache_info()
+      cache_info2 = pjit_lib._pjit_lower_cached.cache_info()
       self.assertIsInstance(out2.sharding, NamedSharding)
 
     # Drops out of C++ cache i.e. cache miss
@@ -3613,11 +3629,13 @@ class ArrayPjitTest(jtu.JaxTestCase):
     f = jax.jit(lambda x: x * 2)
     out = f(arr)
     cache_info1 = pxla._cached_compilation.cache_info()
+    pl_cache_info1 = pjit_lib._pjit_lower_cached.cache_info()
     self.assertIsInstance(out.sharding, NamedSharding)
 
     with jtu.count_pjit_cpp_cache_miss() as count:
       out2 = f(arr2)
       cache_info2 = pxla._cached_compilation.cache_info()
+      pl_cache_info2 = pjit_lib._pjit_lower_cached.cache_info()
       self.assertIsInstance(out2.sharding, PositionalSharding)
 
       # This will hit the cpp cache.
@@ -3628,6 +3646,9 @@ class ArrayPjitTest(jtu.JaxTestCase):
     self.assertEqual(cache_info2.hits, cache_info1.hits + 1)
     self.assertEqual(cache_info2.misses, cache_info1.misses)
 
+    self.assertEqual(pl_cache_info2.hits, pl_cache_info1.hits)
+    self.assertEqual(pl_cache_info2.misses, pl_cache_info1.misses + 1)
+
     out4 = jnp.sum(arr)
     self.assertIsInstance(out4.sharding, NamedSharding)