Merge pull request #4368 from AI-Hypercomputer:darisoy-fix-kvcache-tpu7x

PiperOrigin-RevId: 944600894

Author: Google-ML-Automation

.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/end_to_end/tpu/qwen3/30b/test_qwen3.sh

@@ -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