update unit-test yml

Signed-off-by: Juncheng Gu <jcgu@google.com>

Author: juncgu-google

.buildkite/features/KV_Cache_Offload.yml

@@ -4,8 +4,11 @@ steps:
   - label: "Correctness tests for KV Cache Offload"
     key: "KV_Cache_Offload_CorrectnessTest"
     soft_fail: true
+    env:
+      USE_V6E8_QUEUE: "True"
+      VLLM_LOG_LEVEL: "INFO"
     agents:
-      queue: tpu_v6e_queue
+      queue: tpu_v6e_8_queue
     commands:
       - |
         .buildkite/scripts/run_in_docker.sh \

.buildkite/pipeline_jax.yml

@@ -122,7 +122,7 @@ steps:
            --ignore=/workspace/tpu_inference/tests/e2e \
            --ignore=/workspace/tpu_inference/tpu_inference/mock \
            --ignore=/workspace/tpu_inference/tests/layers/vllm/test_compressed_tensors_moe.py \
-           --ignore=/workspace/tpu_inference/tests/distributed/offload/test_offload_accuracy_test.py \
+           --ignore=/workspace/tpu_inference/tests/distributed/offload \
            --cov-config=/workspace/tpu_inference/.coveragerc --cov tpu_inference --cov-report term-missing --cov-fail-under=69
 
    - label: "JAX unit tests - kernels"
@@ -138,6 +138,7 @@ steps:
              --ignore=/workspace/tpu_inference/tests/kernels/ragged_paged_attention_kernel_v2_test.py \
              --ignore=/workspace/tpu_inference/tests/kernels/ragged_kv_cache_update_v2_test.py \
              --ignore=/workspace/tpu_inference/tests/kernels/collectives \
+             --ignore=/workspace/tpu_inference/tests/kernels/host_dma_test.py \
              --ignore=/workspace/tpu_inference/tests/kernels/fused_moe_v1_test.py
          else
            echo "Skipping: no changes detected in kernels, tests/kernels, or requirements.txt"
@@ -256,6 +257,21 @@ steps:
            echo "Skipping: NIGHTLY environment variable not set"
            exit 0
          fi
+
+   - label: "kv cache offload tests on multi chips"
+     key: test_17
+     soft_fail: true
+     env:
+       USE_V6E8_QUEUE: "True"
+       VLLM_LOG_LEVEL: "INFO"
+     agents:
+       queue: tpu_v6e_8_queue
+     commands:
+       - |
+         .buildkite/scripts/run_in_docker.sh \
+           python3 -m pytest -s -v -x /workspace/tpu_inference/tests/distributed/offload/ \
+           /workspace/tpu_inference/tests/kernels/host_dma_test.py \
+           --ignore=/workspace/tpu_inference/tests/distributed/offload/tpu_offload_accuracy_test.py
   # -----------------------------------------------------------------
   # NOTIFICATION STEP
   # -----------------------------------------------------------------
@@ -278,9 +294,10 @@ steps:
        - test_13
        - test_15
        - test_16
+       - test_17
      agents:
        queue: cpu
      commands:
        - |
          .buildkite/scripts/check_results.sh \
-           "TPU JAX Tests Failed" test_0 test_1 test_2 test_3 test_4 test_5 test_6 test_7 test_8 test_9 test_10 test_11 test_12 test_13 test_15 test_16
+           "TPU JAX Tests Failed" test_0 test_1 test_2 test_3 test_4 test_5 test_6 test_7 test_8 test_9 test_10 test_11 test_12 test_13 test_15 test_16 test_17

tests/distributed/offload/tpu_offload_utils_test.py

@@ -17,7 +17,8 @@ def setUp(self):
         """Set up common parameters for the tests."""
         self.num_layers = 2
         self.num_tokens = 256
-        self.num_kv_heads = 8
+        num_devices = len(list(jax.devices()))
+        self.num_kv_heads = num_devices
         self.head_dim = 128
         self.block_size = 16
         self.num_blocks = self.num_tokens // self.block_size
@@ -37,7 +38,7 @@ def setUp(self):
 
         self.cache_dtype = jnp.bfloat16
 
-        self.mesh = self.create_mesh((1, 8), ("data", "model"))
+        self.mesh = self.create_mesh((1, num_devices), ("data", "model"))
         partition_spec = PartitionSpec(None, None, "model")
         self.device_sharding = NamedSharding(self.mesh,
                                              partition_spec,

tests/kernels/host_dma_test.py

@@ -6,7 +6,6 @@
 import jax.numpy as jnp
 import numpy as np
 from absl.testing import absltest, parameterized
-from jax._src import compilation_cache as cc
 from jax._src import test_util as jtu
 from jax.sharding import NamedSharding, PartitionSpec
 
@@ -15,7 +14,6 @@
 DATA_LOCATION = Literal["device", "host"]
 
 
-# TODO(jcgu): add into CI tests
 @jtu.with_config(jax_numpy_dtype_promotion='strict')
 class HostHbmDmaTest(jtu.JaxTestCase):
 
@@ -27,9 +25,7 @@ def setUp(self):
 
     def tearDown(self):
         super().tearDown()
-        # Reset the cache after each test.
-        # This can also be achieved by running with JAX_TEST_WITH_PERSISTENT_COMPILATION_CACHE=True
-        cc.reset_cache()
+        jax.clear_caches()
 
     def create_mesh(self, axis_shapes, axis_names):
         """Creates a JAX device mesh with the default device order."""