chore: update attention tests and README

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: sagarchapara

README.md

@@ -210,6 +210,7 @@ After installation completes, run the training script.
   --xla_enable_async_all_gather=true \
   --xla_tpu_scoped_vmem_limit_kib=65536 \
   --xla_tpu_enable_async_all_to_all=true \
+  --xla_tpu_enable_latency_hiding_scheduler=true \
   --xla_tpu_enable_all_experimental_scheduler_features=true \
   --xla_tpu_enable_scheduler_memory_pressure_tracking=true \
   --xla_tpu_host_transfer_overlap_limit=24 \
@@ -329,6 +330,7 @@ After installation completes, run the training script.
   --xla_enable_async_all_gather=true \
   --xla_tpu_scoped_vmem_limit_kib=65536 \
   --xla_tpu_enable_async_all_to_all=true \
+  --xla_tpu_enable_latency_hiding_scheduler=true \
   --xla_tpu_enable_all_experimental_scheduler_features=true \
   --xla_tpu_enable_scheduler_memory_pressure_tracking=true \
   --xla_tpu_host_transfer_overlap_limit=24 \

src/maxdiffusion/tests/attention_test.py

@@ -45,7 +45,7 @@ def _ulysses_mesh(self):
 
   def _ulysses_ring_mesh(self):
     devices = np.array(jax.devices()[:4]).reshape(1, 1, 2, 2)
-    return Mesh(devices, ("data", "fsdp", "context", "tensor"))
+    return Mesh(devices, ("data", "fsdp", "ring", "ulysses"))
 
   def _ulysses_axis_rules(self):
     return (
@@ -60,8 +60,8 @@ def _ulysses_ring_axis_rules(self):
     return (
         (attention_flax.BATCH, "data"),
         (attention_flax.SELF_ATTN_HEAD, None),
-        (attention_flax.SELF_ATTN_Q_LENGTH, ("context", "tensor")),
-        (attention_flax.SELF_ATTN_KV_LENGTH, ("context", "tensor")),
+        (attention_flax.SELF_ATTN_Q_LENGTH, ("ring", "ulysses")),
+        (attention_flax.SELF_ATTN_KV_LENGTH, ("ring", "ulysses")),
         (attention_flax.D_KV, None),
     )