staging_tokens --> staging_blocks

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

Author: juncgu-google

tests/distributed/offload/tpu_offload_connector_scheduler_test.py

@@ -61,15 +61,15 @@ def scheduler_factory():
     def _scheduler(
         block_size: int = _DEFAULT_BLOCK_SIZE,
         offload_decode_save: int = 0,
-        offload_staging_buffer_tokens: int = -1,
+        offload_num_staging_blocks: int = -1,
         offload_num_cpu_chunks: int = -1,
     ):
         # update config
         vllm_config = MockVllmConfig(block_size=block_size)
         os.environ["TPU_OFFLOAD_DECODE_SAVE"] = str(offload_decode_save)
-        if offload_staging_buffer_tokens >= 0:
-            os.environ["TPU_OFFLOAD_STAGING_BUFFER_TOKENS"] = str(
-                offload_staging_buffer_tokens)
+        if offload_num_staging_blocks >= 0:
+            os.environ["TPU_OFFLOAD_NUM_STAGING_BLOCKS"] = str(
+                offload_num_staging_blocks)
         if offload_num_cpu_chunks > 0:
             os.environ["TPU_OFFLOAD_NUM_CPU_CHUNKS"] = str(
                 offload_num_cpu_chunks)
@@ -111,9 +111,8 @@ def test_get_num_new_matched_tokens_hit(self, scheduler_factory,
         5. skip 1 block + full-hit + only 1 staging block
         6. skip 1 block + full-hit + no staging block
         """
-        num_staging_tokens = num_staging_blocks * _DEFAULT_BLOCK_SIZE
         scheduler = scheduler_factory(
-            offload_staging_buffer_tokens=num_staging_tokens)
+            offload_num_staging_blocks=num_staging_blocks)
         prompt_len = scheduler.block_size * num_prompt_blocks
         num_computed_tokens = scheduler.block_size * num_computed_blocks
         num_blocks_to_load = num_matched_blocks - num_computed_blocks
@@ -231,7 +230,7 @@ def test_build_connector_meta_new_prefill(self, scheduler_factory,
         """
         num_staging_blocks = num_staging_tokens // _DEFAULT_BLOCK_SIZE
         scheduler = scheduler_factory(
-            offload_staging_buffer_tokens=num_staging_tokens,
+            offload_num_staging_blocks=num_staging_blocks,
             offload_num_cpu_chunks=100)
 
         # calculate the groundtruth
@@ -347,10 +346,9 @@ def test_build_connector_meta_decode_with_save(self, scheduler_factory,
         2. th N-th decode (hit block bounary) + not decode_save (no save)
         """
 
-        scheduler = scheduler_factory(
-            offload_decode_save=decode_save,
-            offload_staging_buffer_tokens=_DEFAULT_BLOCK_SIZE * 10,
-            offload_num_cpu_chunks=10)
+        scheduler = scheduler_factory(offload_decode_save=decode_save,
+                                      offload_num_staging_blocks=10,
+                                      offload_num_cpu_chunks=10)
 
         prompt_tokens = list(range(prompt_len))
         generated_tokens = list(range(prompt_len, seq_len))

tpu_inference/distributed/offload/tpu_offload_connector.py

@@ -511,9 +511,8 @@ def __init__(self, vllm_config: "VllmConfig"):
 
         # config staging buffer
         # NOTE(jcgu): Need to find a way to grab page_size_bytes in scheduler
-        # otherwise, we can only use # of tokens as input, instead of buffer size in GB
-        num_staging_buffer_tokens = envs.TPU_OFFLOAD_STAGING_BUFFER_TOKENS
-        self.num_staging_blocks = num_staging_buffer_tokens // self.block_size
+        # otherwise, we can only use # of blocks as input, instead of buffer size in GB
+        self.num_staging_blocks = envs.TPU_OFFLOAD_NUM_STAGING_BLOCKS
         self.staging_buffer_manager = StagingBufferManager(
             num_blocks=self.num_staging_blocks)
 
@@ -698,19 +697,15 @@ def _prepare_req_meta(
         block_hashes = self._get_request_block_hashes(_request)
         self.offload_manager.touch(block_hashes)
 
-        # only consider the tokens covered by block_hashes
+        # only consider the tokens covered by block_hashes;
+        # currently full blocks only
         num_total_blocks = len(block_hashes)
         num_total_tokens = min(num_total_blocks * self.block_size,
                                len(tracker.token_ids))
         num_full_blocks = num_total_tokens // self.block_size
-        num_full_blocks_tokens = num_full_blocks * self.block_size
-        # adjust last partial block
-        last_partial_block_num_tokens = num_total_tokens - num_full_blocks_tokens
-        need_last_block = self._adjust_last_partial_block(
-            last_partial_block_num_tokens)
-        adjusted_num_total_tokens = num_total_tokens if need_last_block else num_full_blocks_tokens
-        adjusted_num_total_blocks = num_full_blocks + (1 if need_last_block
-                                                       else 0)
+        num_full_block_tokens = num_full_blocks * self.block_size
+        adjusted_num_total_tokens = num_full_block_tokens
+        adjusted_num_total_blocks = num_full_blocks
         assert adjusted_num_total_blocks <= len(tracker.block_ids)
 
         has_new_tokens = adjusted_num_total_tokens > tracker.save_watermark

tpu_inference/envs.py

@@ -28,7 +28,7 @@
     TPU_OFFLOAD_SWAP_OP_TYPE: str = "jax"
     TPU_OFFLOAD_DECODE_SAVE: bool = False
     TPU_OFFLOAD_NUM_CPU_CHUNKS: int = 1024
-    TPU_OFFLOAD_STAGING_BUFFER_TOKENS: int = 8192
+    TPU_OFFLOAD_NUM_STAGING_BLOCKS: int = 128
 
 
 def env_with_choices(
@@ -127,21 +127,21 @@ def _get_validated_env() -> str | None:
     # Ray compiled DAG channel type for TPU
     "VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE":
     env_with_choices("VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE", "shm", ["shm"]),
-    # kv offload to dram: save kv in the decode phase
-    "TPU_OFFLOAD_DECODE_SAVE":
-    lambda: bool(int(os.getenv("TPU_OFFLOAD_DECODE_SAVE", "0"))),
+    # kv offload to dram: skip pre-compiling swap-related jax functions
+    "TPU_OFFLOAD_SKIP_JAX_PRECOMPILE":
+    lambda: bool(int(os.getenv("TPU_OFFLOAD_SKIP_JAX_PRECOMPILE", "0"))),
     # kv offload to dram: swap function type: jax, or pallas
     "TPU_OFFLOAD_SWAP_OP_TYPE":
     lambda: os.getenv("TPU_OFFLOAD_SWAP_OP_TYPE", "jax"),
+    # kv offload to dram: save kv in the decode phase
+    "TPU_OFFLOAD_DECODE_SAVE":
+    lambda: bool(int(os.getenv("TPU_OFFLOAD_DECODE_SAVE", "0"))),
     # kv offload to dram: dram space size in # of chunks / blocks
     "TPU_OFFLOAD_NUM_CPU_CHUNKS":
     lambda: int(os.getenv("TPU_OFFLOAD_NUM_CPU_CHUNKS", "1024")),
-    # kv offload to dram: dram space size in # of chunks / blocks
-    "TPU_OFFLOAD_SKIP_JAX_PRECOMPILE":
-    lambda: bool(int(os.getenv("TPU_OFFLOAD_SKIP_JAX_PRECOMPILE", "0"))),
     # kv offload to dram: size of staging buffer (hbm) for swap
-    "TPU_OFFLOAD_STAGING_BUFFER_TOKENS":
-    lambda: int(os.getenv("TPU_OFFLOAD_STAGING_BUFFER_TOKENS", "16384")),
+    "TPU_OFFLOAD_NUM_STAGING_BLOCKS":
+    lambda: int(os.getenv("TPU_OFFLOAD_NUM_STAGING_BLOCKS", "128")),
 }
 
 

tpu_inference/worker/tpu_worker.py

@@ -294,9 +294,7 @@ def determine_available_memory(self) -> int:
             kv_transfer_config = self.vllm_config.kv_transfer_config
             if kv_transfer_config.kv_connector == "TPUOffloadConnector" and kv_transfer_config.kv_connector_module_path == "tpu_inference.distributed.offload.tpu_offload_connector":
                 # If kv offloading is enabled, we need to account for the memory used by the KV transfer buffer.
-                staging_buffer_tokens = envs.TPU_OFFLOAD_STAGING_BUFFER_TOKENS
-                # calculate staging buffer size
-                staging_buffer_pages = staging_buffer_tokens // self.vllm_config.cache_config.block_size
+                staging_buffer_pages = envs.TPU_OFFLOAD_NUM_STAGING_BLOCKS
 
                 kv_cache_specs = self.model_runner.get_kv_cache_spec()
                 num_layers = len(kv_cache_specs)