Merge pull request #3404 from ROCm:mla-test-ici-parallelism

PiperOrigin-RevId: 886875950

Author: Google-ML-Automation

tests/unit/attention_test.py

@@ -1611,6 +1611,7 @@ def get_causal_mask_for_indexer(self, batch_size, q_len, kv_len):
   def test_indexer_loss(self):
     """Test indexer loss computation."""
     mla_config_args = self.config_arguments.copy()
+    mla_config_args.update(get_decoupled_parallelism_overrides())
     mla_config_args["use_sparse_indexer"] = True
     mla_config_args["attention"] = "dot_product"
     _, mla = self.init_mla(mla_config_args, rope_type="default")
@@ -1657,6 +1658,7 @@ def test_indexer_loss(self):
   def test_indexer_loss_kl_divergence_zero(self):
     """Test that KL divergence is 0 when target and pred distributions match exactly."""
     mla_config_args = self.config_arguments.copy()
+    mla_config_args.update(get_decoupled_parallelism_overrides())
     mla_config_args["use_sparse_indexer"] = True
     mla_config_args["attention"] = "dot_product"
     _, mla = self.init_mla(mla_config_args, rope_type="default")