update ci tests

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

Author: juncgu-google

.buildkite/features/KV_Cache_Host_Offloading.yml

@@ -9,7 +9,7 @@ steps:
     commands:
       - |
         .buildkite/scripts/run_in_docker.sh \
-          python3 -m pytest -s -v /workspace/tpu_inference/tests/distributed/offload/
+          python3 -m pytest -s -v /workspace/tpu_inference/tests/distributed/offload/tpu_offload_accuracy_test.py
   - label: "Record correctness test result for KV Cache Offload"
     key: "record_KV_Cache_Offload_CorrectnessTest"
     depends_on: "KV_Cache_Offload_CorrectnessTest"

.buildkite/features/KV_Cache_Offload.yml

@@ -122,6 +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 \
            --cov-config=/workspace/tpu_inference/.coveragerc --cov tpu_inference --cov-report term-missing --cov-fail-under=69
 
    - label: "JAX unit tests - kernels"

.buildkite/pipeline_jax.yml

@@ -60,7 +60,6 @@ def _test_kv_cache_cpu_offloading_accuracy(
         os.environ['TPU_OFFLOAD_DECODE_SAVE'] = decode_save
         llm = LLM(model="meta-llama/Llama-3.2-3B",
                   max_model_len=1024,
-                  tensor_parallel_size=8,
                   task="generate",
                   kv_transfer_config=kv_transfer_config)
 

tests/distributed/offload/tpu_offload_accuracy_test.py

@@ -72,9 +72,10 @@ def setUp(self):
         self.num_blocks = 128
         self.num_cpu_chunks = 128
         self.block_size = self.vllm_config.cache_config.block_size
-        self.num_heads = 8
+        num_devices = len(list(jax.devices()))
+        self.num_heads = num_devices
         self.head_size = 128
-        self.mesh = self.create_mesh((1, 8), ("data", "model"))
+        self.mesh = self.create_mesh((1, num_devices), ("data", "model"))
         if self.mesh is None:
             self.skipTest("Cannot create mesh. Must be run on a TPU node.")
             return