[Mosaic GPU] Avoid double-predication when async_copy predicate is specified

PiperOrigin-RevId: 700999181

Author: apaszke

jax/experimental/mosaic/gpu/core.py

@@ -356,7 +356,7 @@ def async_copy(
       arrive: bool | None = None,
       uniform: bool = True,
       collective: Sequence[gpu.Dimension] | gpu.Dimension | None = None,
-      predicate: ir.Value | None = None,
+      predicate: ir.Value | None = None,  # Should select 0 or 1 threads from the WG.
   ):
     index = ir.IndexType.get()
     i16 = ir.IntegerType.get_signless(16)
@@ -504,7 +504,7 @@ def partition_dim(dim: int, idx: ir.Value, num_chunks: int):
 
     uniform_ctx = (
         functools.partial(utils.single_thread, per_block=False)
-        if uniform
+        if uniform and predicate is None
         else contextlib.nullcontext
     )