offload envs

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

Author: juncgu-google

tests/distributed/offload/tpu_offload_connector_scheduler_test.py

@@ -62,18 +62,12 @@ def scheduler_factory():
     def _scheduler(
         block_size: int = _DEFAULT_BLOCK_SIZE,
         offload_decode_save: int = 0,
-        offload_partial_block_save_behavior: str = "drop",
-        offload_partial_block_dynamic_pad_lower_limit: int = 0,
         offload_staging_buffer_tokens: int = -1,
         offload_num_cpu_chunks: int = DEFAULT_TPU_OFFLOAD_CPU_CHUNKS,
     ):
         # update config
         vllm_config = MockVllmConfig(block_size=block_size)
         os.environ["TPU_OFFLOAD_DECODE_SAVE"] = str(offload_decode_save)
-        os.environ[
-            "TPU_OFFLOAD_PARTIAL_BLOCK_SAVE_BEHAVIOR"] = offload_partial_block_save_behavior
-        os.environ["TPU_OFFLOAD_PARTIAL_BLOCK_DYNAMIC_PAD_LOWER_LIMIT"] = str(
-            offload_partial_block_dynamic_pad_lower_limit)
         if offload_staging_buffer_tokens >= 0:
             os.environ["TPU_OFFLOAD_STAGING_BUFFER_TOKENS"] = str(
                 offload_staging_buffer_tokens)
@@ -238,7 +232,6 @@ def test_build_connector_meta_new_prefill(self, scheduler_factory,
         """
         num_staging_blocks = num_staging_tokens // _DEFAULT_BLOCK_SIZE
         scheduler = scheduler_factory(
-            offload_partial_block_save_behavior="drop",
             offload_staging_buffer_tokens=num_staging_tokens,
             offload_num_cpu_chunks=100)
 

tpu_inference/distributed/offload/cpu_backend.py

@@ -10,9 +10,6 @@
 
 logger = init_logger(__name__)
 
-GB = 1024**3
-DEFAULT_CPU_CACHE_SIZE_BYTES = 1 * GB
-
 
 class LocalCPUBackend:
     """

tpu_inference/distributed/offload/offload_manager.py

@@ -12,9 +12,6 @@
 
 logger = init_logger(__name__)
 
-GB = 1024**3
-DEFAULT_CPU_CACHE_SIZE_BYTES = 1 * GB
-
 ChunkHash = BlockHash
 
 

tpu_inference/distributed/offload/tpu_offload_connector.py

@@ -112,13 +112,13 @@
     from vllm.v1.request import Request
     from vllm.forward_context import ForwardContext
 
+from tpu_inference import envs
 from tpu_inference.distributed.offload.cpu_backend import LocalCPUBackend
 from tpu_inference.distributed.offload.offload_manager import (
     LRUCacheManager, StagingBufferManager)
 from tpu_inference.distributed.offload.utils import (
     CPU_OFFLOADING_SWAP_OP_TYPE, CpuChunkId, KVCacheSwapFn, ReqId,
-    TokenProcessor, get_default_kv_connector_staging_buffer_tokens,
-    get_kv_cache_swap_fn, jitted_insert_kv_cache_slices)
+    TokenProcessor, 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
@@ -480,9 +480,7 @@ def __init__(self, vllm_config: "VllmConfig"):
         self.block_size = vllm_config.cache_config.block_size
 
         # offloading manager
-        self.num_cpu_chunks = int(
-            os.getenv("TPU_OFFLOAD_NUM_CPU_CHUNKS",
-                      str(DEFAULT_TPU_OFFLOAD_CPU_CHUNKS)))
+        self.num_cpu_chunks = envs.TPU_OFFLOAD_NUM_CPU_CHUNKS
         self.offload_manager = LRUCacheManager(
             num_cpu_chunks=self.num_cpu_chunks)
 
@@ -506,14 +504,15 @@ def __init__(self, vllm_config: "VllmConfig"):
         self.token_processor = TokenProcessor(model_name=model_name,
                                               chunk_size=self.block_size)
 
-        self.decode_save = os.getenv("TPU_OFFLOAD_DECODE_SAVE", "0") == "1"
+        self.decode_save = envs.TPU_OFFLOAD_DECODE_SAVE
         # NOTE(jcgu): currently, let's make chunk_size == block_size
         # chunk_size == n * block_size lead to
         #  1. multi-size chunks
         #  2. complicated resize (split, concatenate) operations due to
         #     real-chunk-size in save and load
         self.cpu_chunk_size = self.block_size
 
+        # TODO(jcgu): rm
         # define partial_block saving behavior
         self.partial_block_save_behavior: PARTIAL_BLOCK_SAVE_BEHAVIOR = \
             os.getenv("TPU_OFFLOAD_PARTIAL_BLOCK_SAVE_BEHAVIOR", "drop")
@@ -535,11 +534,7 @@ 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
-        _default_staging_buffer_tokens = get_default_kv_connector_staging_buffer_tokens(
-        )
-        num_staging_buffer_tokens = int(
-            os.getenv("TPU_OFFLOAD_STAGING_BUFFER_TOKENS",
-                      str(_default_staging_buffer_tokens)))
+        num_staging_buffer_tokens = envs.TPU_OFFLOAD_STAGING_BUFFER_TOKENS
         self.num_staging_blocks = num_staging_buffer_tokens // self.block_size
         self.staging_buffer_manager = StagingBufferManager(
             num_blocks=self.num_staging_blocks)
@@ -1214,15 +1209,13 @@ def __init__(self, vllm_config: VllmConfig,
 
         self.runner: Optional[TPUModelRunner] = None
         self.mesh: Optional[Mesh] = None
-        self.swap_op_type = os.getenv("TPU_OFFLOAD_SWAP_OP_TYPE",
-                                      default=DEFAULT_HOST_HBM_SWAP_OP_TYPE)
+        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}")
 
-        self.use_bucketed_swap_ops = os.getenv(
-            "TPU_OFFLOAD_SKIP_JAX_PRECOMPILE", "0") == "0"
+        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}"
         )
@@ -1231,9 +1224,7 @@ def __init__(self, vllm_config: VllmConfig,
         self.swap_out_fn: KVCacheSwapFn = None
 
         # cpu cache
-        self.num_cpu_chunks = int(
-            os.getenv("TPU_OFFLOAD_NUM_CPU_CHUNKS",
-                      str(DEFAULT_TPU_OFFLOAD_CPU_CHUNKS)))
+        self.num_cpu_chunks = envs.TPU_OFFLOAD_NUM_CPU_CHUNKS
         self.cpu_backend = LocalCPUBackend(num_cpu_chunks=self.num_cpu_chunks)
         # The worker needs its own token processor to generate keys.
         model_name = self.vllm_config.model_config.model

tpu_inference/envs.py

@@ -24,6 +24,11 @@
     NUM_SLICES: int = 1
     RAY_USAGE_STATS_ENABLED: str = "0"
     VLLM_USE_RAY_COMPILED_DAG_CHANNEL_TYPE: str = "shm"
+    TPU_OFFLOAD_SKIP_JAX_PRECOMPILE: bool = False
+    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
 
 
 def env_with_choices(
@@ -122,6 +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: swap function type: jax, or pallas
+    "TPU_OFFLOAD_SWAP_OP_TYPE":
+    lambda: os.getenv("TPU_OFFLOAD_SWAP_OP_TYPE", "jax"),
+    # 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")),
 }