Merge pull request #3041 from Steboss:main

PiperOrigin-RevId: 864463680

Author: Google-ML-Automation

src/MaxText/layers/moe.py

@@ -950,9 +950,13 @@ def gmm(inputs, kernel, tiling, group_sizes, expert_assignments, weight_gather_a
             # Use full contraction for QWIX quantization to allow quantization
             # fusion (max reduce over contracting dimension).
             tiling = (tiling[0], k, tiling[2])
+
+          is_tpu = (self.mesh.devices.flat[0] == "tpu")
+          # TPU needs random mosaic_fusion_group; GPU/CPU needs deterministic ID for autotuner sync
+          mosaic_group_id = f"{random.randint(0, 1000000000)}" if is_tpu else '0'
           with set_xla_metadata(
               ragged_dot_tiling=",".join([str(t) for t in tiling]),
-              mosaic_fusion_group=f"{random.randint(0, 1000000000)}",
+              mosaic_fusion_group=mosaic_group_id,
           ):
             output = jax.lax.ragged_dot(
                 lhs=inputs,