worker_test: multi requests; acc_test: precompile

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

Author: juncgu-google

tests/distributed/offload/tpu_offload_accuracy_test.py

@@ -1,5 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 
+import itertools
 import os
 import time
 
@@ -49,12 +50,13 @@ def _test_kv_cache_cpu_offloading_accuracy(
     sampling_config: SamplingParams,
     kv_transfer_config: KVTransferConfig,
     swap_op_type: str,
+    skip_precompile: str,
     decode_save: str,
 ):
     with monkeypatch.context():
         os.environ['SKIP_JAX_PRECOMPILE'] = '1'
-        os.environ['TPU_OFFLOAD_SKIP_JAX_PRECOMPILE'] = '1'
         os.environ['TPU_OFFLOAD_SWAP_OP_TYPE'] = swap_op_type
+        os.environ['TPU_OFFLOAD_SKIP_JAX_PRECOMPILE'] = skip_precompile
         os.environ['TPU_OFFLOAD_DECODE_SAVE'] = decode_save
         llm = LLM(model="meta-llama/Llama-3.2-3B",
                   max_model_len=1024,
@@ -98,12 +100,14 @@ def test_kv_cache_cpu_offloading_accuracy(
 ):
     swap_op_types = ["pallas", "jax"]
     decode_saves = ["0", "1"]
-    for swap_op_type in swap_op_types:
-        for decode_save in decode_saves:
-            _test_kv_cache_cpu_offloading_accuracy(
-                monkeypatch,
-                sampling_config,
-                kv_transfer_config,
-                swap_op_type,
-                decode_save,
-            )
+    skip_precompile = ["0", "1"]
+    for swap_op_type, decode_save, _skip_precompile in itertools.product(
+            swap_op_types, decode_saves, skip_precompile):
+        _test_kv_cache_cpu_offloading_accuracy(
+            monkeypatch,
+            sampling_config,
+            kv_transfer_config,
+            swap_op_type,
+            _skip_precompile,
+            decode_save,
+        )

tests/distributed/offload/tpu_offload_connector_worker_test.py

@@ -70,7 +70,7 @@ def setUp(self):
         self.vllm_config = MockVllmConfig(block_size=_DEFAULT_BLOCK_SIZE)
         self.num_layers = 80
         self.num_blocks = 128
-        self.num_cpu_chunks = 24
+        self.num_cpu_chunks = 128
         self.block_size = self.vllm_config.cache_config.block_size
         self.num_heads = 8
         self.head_size = 128
@@ -205,40 +205,57 @@ def test_precompile_run_success(self, swap_op_type: str):
 
     @parameterized.named_parameters(
         dict(
-            testcase_name="_regular_single_block_save",
+            testcase_name="_single_block",
             num_blocks_to_save=1,
             num_requests=1,
         ),
         dict(
-            testcase_name="_regular_multi_requests_single_block_save",
-            num_blocks_to_save=2,
-            num_requests=4,
+            testcase_name="_multi_requests_single_block",
+            num_blocks_to_save=1,
+            num_requests=6,
         ),
         dict(
-            testcase_name="_regular_multi_block_save",
+            testcase_name="_multi_blocks",
             num_blocks_to_save=5,
             num_requests=1,
         ),
         dict(
-            testcase_name="_regular_multi_block_save_with_compile_jax",
+            testcase_name="_multi_requests_multi_blocks",
+            num_blocks_to_save=5,
+            num_requests=6,
+        ),
+        dict(
+            testcase_name="_multi_blocks_with_compile_jax",
             num_blocks_to_save=5,
             num_requests=1,
             use_precompiled_swap_ops=True,
         ),
         dict(
-            testcase_name=
-            "_regular_multi_request_single_block_save_with_compile_jax",
+            testcase_name="_multi_requests_single_block_with_compile_jax",
             num_blocks_to_save=1,
             num_requests=6,
             use_precompiled_swap_ops=True,
         ),
         dict(
-            testcase_name="_regular_multi_block_save_with_compile_pallas",
+            testcase_name="_multi_requests_multi_blocks_with_compile_jax",
+            num_blocks_to_save=5,
+            num_requests=6,
+            use_precompiled_swap_ops=True,
+        ),
+        dict(
+            testcase_name="_multi_blocks_with_compile_pallas",
             num_blocks_to_save=5,
             num_requests=1,
             use_precompiled_swap_ops=True,
             swap_op_type="pallas",
         ),
+        dict(
+            testcase_name="_multi_requests_multi_blocks_with_compile_pallas",
+            num_blocks_to_save=5,
+            num_requests=6,
+            use_precompiled_swap_ops=True,
+            swap_op_type="pallas",
+        ),
         dict(
             testcase_name="_final_save",
             num_blocks_to_save=1,
@@ -370,13 +387,13 @@ def test_tpu_connector_save(
 
     @parameterized.named_parameters(
         dict(
-            testcase_name="_single_block_",
+            testcase_name="_single_block",
             num_blocks_to_operate=1,
             num_requests=1,
         ),
         dict(
-            testcase_name="_multi_requests_",
-            num_blocks_to_operate=2,
+            testcase_name="_multi_requests_single_block",
+            num_blocks_to_operate=1,
             num_requests=4,
         ),
         dict(
@@ -387,9 +404,23 @@ def test_tpu_connector_save(
             swap_op_type="jax",
         ),
         dict(
-            testcase_name="_multi_blocks_compile_pallas",
+            testcase_name="_multi_requests_single_block_compile_jax",
+            num_blocks_to_operate=1,
+            num_requests=6,
+            use_precompiled_swap_ops=True,
+            swap_op_type="jax",
+        ),
+        dict(
+            testcase_name="_multi_requests_multi_blocks_compile_jax",
             num_blocks_to_operate=5,
-            num_requests=1,
+            num_requests=6,
+            use_precompiled_swap_ops=True,
+            swap_op_type="jax",
+        ),
+        dict(
+            testcase_name="_multi_requests_multi_blocks_compile_pallas",
+            num_blocks_to_operate=5,
+            num_requests=6,
             use_precompiled_swap_ops=True,
             swap_op_type="pallas",
         ),

tpu_inference/distributed/offload/tpu_offload_connector.py

@@ -118,7 +118,7 @@
     LRUCacheManager, StagingBufferManager)
 from tpu_inference.distributed.offload.utils import (
     CPU_OFFLOADING_SWAP_OP_TYPE, CpuChunkId, KVCacheSwapFn, ReqId,
-    TokenProcessor, get_kv_cache_swap_fn, jitted_insert_kv_cache_slices)
+    get_kv_cache_swap_fn, jitted_insert_kv_cache_slices)
 from tpu_inference.logger import init_logger
 from tpu_inference.runner.kv_cache_manager import KVCacheManager
 from tpu_inference.runner.tpu_runner import TPUModelRunner
@@ -496,8 +496,6 @@ def __init__(self, vllm_config: "VllmConfig"):
         self._reqs_being_loaded = defaultdict[ReqId, set[CpuChunkId]](set)
 
         model_name = self.vllm_config.model_config.model
-        self.token_processor = TokenProcessor(model_name=model_name,
-                                              chunk_size=self.block_size)
 
         self.decode_save = envs.TPU_OFFLOAD_DECODE_SAVE
         # NOTE(jcgu): currently, let's make chunk_size == block_size
@@ -528,7 +526,7 @@ def __init__(self, vllm_config: "VllmConfig"):
 
     def _get_request_block_hashes(self, req: "Request") -> list[BlockHash]:
         # request's original block_hashes do not include the last partial block
-        # TODO(jcgu): switch back to token_processor
+        # TODO(jcgu): add an option to use local token_processor
         return req.block_hashes
 
     def get_num_new_matched_tokens(
@@ -1160,19 +1158,14 @@ def __init__(self, vllm_config: VllmConfig,
 
         self.runner: Optional[TPUModelRunner] = None
         self.mesh: Optional[Mesh] = None
+        self.swap_in_fn: KVCacheSwapFn = None
+        self.swap_out_fn: KVCacheSwapFn = None
         self.swap_op_type = envs.TPU_OFFLOAD_SWAP_OP_TYPE
-        assert self.swap_op_type in get_args(CPU_OFFLOADING_SWAP_OP_TYPE)
         # TODO(jcgu): check libtpu compatibility for pallas dma kernel
-        logger.info(
-            f"(cpu offloading) swap operation type is {self.swap_op_type}")
-
+        assert self.swap_op_type in get_args(CPU_OFFLOADING_SWAP_OP_TYPE)
         self.use_bucketed_swap_ops = not envs.TPU_OFFLOAD_SKIP_JAX_PRECOMPILE
-        logger.info(
-            f"(cpu offloading) use_bucketed_swap_ops={self.use_bucketed_swap_ops}"
-        )
-
-        self.swap_in_fn: KVCacheSwapFn = None
-        self.swap_out_fn: KVCacheSwapFn = None
+        logger.info(f" swap operation type is {self.swap_op_type}, "
+                    f"use_bucketed_swap_ops={self.use_bucketed_swap_ops}.")
 
         # cpu cache
         self.num_cpu_chunks = envs.TPU_OFFLOAD_NUM_CPU_CHUNKS
@@ -1181,13 +1174,11 @@ def __init__(self, vllm_config: VllmConfig,
         model_name = self.vllm_config.model_config.model
         logger.info(
             f"Model name is {model_name}, KV block_size={self.block_size}")
-        self.token_processor = TokenProcessor(model_name=model_name,
-                                              chunk_size=self.block_size)
 
         self.cpu_chunk_size = self.block_size
         # Thread pool for asynchronous TPU->CPU copies
-        self.save_executor = ThreadPoolExecutor(max_workers=4,
-                                                thread_name_prefix="tpu_saver")
+        self.save_executor = ThreadPoolExecutor(
+            max_workers=4, thread_name_prefix="tpu_save_handler")
         self.finished_save_reqs: set[ReqId] = set()
         self.finished_load_reqs: set[ReqId] = set()
         # Tracks if wait_for_save has been called for the current step's metadata.
@@ -1298,10 +1289,11 @@ def _precompile_kv_swap_operations(self):
 
                 # 3. Pre-compile CPU -> TPU transfer (used in load)
                 split_size_list = [self.block_size] * num_blocks
-                chunked_dummy_kv_cpu = [
-                    jax.lax.split(flat_layer_cache, split_size_list, axis=0)
-                    for flat_layer_cache in dummy_kv_cpu
-                ]
+                chunked_dummy_kv_cpu = jax.tree.map(
+                    lambda flat_layer_cache: jax.lax.split(
+                        flat_layer_cache, split_size_list, axis=0),
+                    dummy_kv_cpu)
+
                 chunked_dummy_kv_tpu = self.swap_in_fn(chunked_dummy_kv_cpu)
                 jax.block_until_ready(chunked_dummy_kv_tpu)
 
@@ -1374,13 +1366,13 @@ def _bucketed_swap_out_fn(
 
         # Fast path: handle bucket-sized transfers
         if num_blocks in BLOCK_SIZE_BUCKETS:
+            split_size_list = [self.block_size] * num_blocks
             flat_kv_caches_cpu = self.swap_out_fn(flat_kv_caches_tpu)
             jax.block_until_ready(flat_kv_caches_cpu)
-            split_size_list = [self.block_size] * num_blocks
-            return [
-                jax.lax.split(flat_layer_cache, split_size_list, axis=0)
-                for flat_layer_cache in flat_kv_caches_cpu
-            ]
+            return jax.tree.map(
+                lambda flat_layer_cache: jax.lax.split(
+                    flat_layer_cache, split_size_list, axis=0),
+                flat_kv_caches_cpu)
 
         # Bucket decomposition path
         decomposed_block_sizes = self._decompose_into_buckets(num_blocks)
@@ -1580,12 +1572,10 @@ def _save_blocks_to_cpu(self, req_id: ReqId, full_block_ids: list[int],
                     # NOTE(jcgu): we keep cpu_chunk_size == block_size
                     split_size_list = [self.cpu_chunk_size
                                        ] * num_blocks_to_save
-                    chunks_on_cpu = [
-                        jax.lax.split(flat_layer_cache,
-                                      split_size_list,
-                                      axis=0)
-                        for flat_layer_cache in flat_kv_caches_cpu
-                    ]
+                    chunks_on_cpu = jax.tree.map(
+                        lambda flat_layer_cache: jax.lax.split(
+                            flat_layer_cache, split_size_list, axis=0),
+                        flat_kv_caches_cpu)
 
             if chunks_on_cpu and chunks_on_cpu[0]:
                 jax.block_until_ready(chunks_on_cpu)