Fix XLA SPMD partitioning hang for kvcache_test.py on 8-device runners and re-enable in TPU7X unit CI

Author: darisoy

.github/workflows/run_tests_coordinator.yml

@@ -195,7 +195,7 @@ jobs:
             "tpu-integration": "--ignore=tests/post_training",
             "tpu-post-training-unit": "",
             "tpu-post-training-integration": "",
-            "tpu7x-unit": "--ignore=tests/post_training --ignore=tests/inference/kvcache_test.py",
+            "tpu7x-unit": "--ignore=tests/post_training",
             "tpu7x-integration": "--ignore=tests/post_training",
             "tpu7x-post-training-unit": "",
             "tpu7x-post-training-integration": "",

tests/inference/kvcache_test.py

@@ -29,7 +29,7 @@ class MlaKVCacheTest(unittest.TestCase):
   def setUp(self):
     super().setUp()
     self.rng = jax.random.PRNGKey(42)
-    self.batchsize = 8
+    self.batchsize = 4
     self.prefill_len = 100
     self.target_len = 196
     self.dtype = jnp.bfloat16