put offload folder under tpu_inference

juncgu-google · juncgu-google · commit 7d81d9010444 · 2025-12-12T07:01:01.000Z
Signed-off-by: Juncheng Gu &lt;jcgu@google.com&gt;
diff --git a/.buildkite/features/KV_Cache_Offload.yml b/.buildkite/features/KV_Cache_Offload.yml
@@ -12,7 +12,7 @@ steps:
     commands:
       - |
         .buildkite/scripts/run_in_docker.sh \
-          python3 -m pytest -s -v /workspace/tpu_inference/tests/distributed/offload/tpu_offload_accuracy_test.py
+          python3 -m pytest -s -v /workspace/tpu_inference/tests/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"
diff --git a/.buildkite/pipeline_jax.yml b/.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 \
+           --ignore=/workspace/tpu_inference/tests/offload \
            --cov-config=/workspace/tpu_inference/.coveragerc --cov tpu_inference --cov-report term-missing --cov-fail-under=69
 
    - label: "JAX unit tests - kernels"
@@ -269,9 +269,9 @@ steps:
      commands:
        - |
          .buildkite/scripts/run_in_docker.sh \
-           python3 -m pytest -s -v -x /workspace/tpu_inference/tests/distributed/offload/ \
+           python3 -m pytest -s -v -x /workspace/tpu_inference/tests/offload/ \
            /workspace/tpu_inference/tests/kernels/host_dma_test.py \
-           --ignore=/workspace/tpu_inference/tests/distributed/offload/tpu_offload_accuracy_test.py
+           --ignore=/workspace/tpu_inference/tests/offload/tpu_offload_accuracy_test.py
   # -----------------------------------------------------------------
   # NOTIFICATION STEP
   # -----------------------------------------------------------------
diff --git a/examples/offload/gke/benchmarks/deploy-cpu-offload.yaml b/examples/offload/gke/benchmarks/deploy-cpu-offload.yaml
@@ -29,7 +29,7 @@ spec:
         imagePullPolicy: Always
         command: ["/bin/sh", "-c"]
         args:
-        - "vllm serve meta-llama/Llama-3.3-70B-Instruct --kv-transfer-config '{\"kv_connector\":\"TPUOffloadConnector\",\"kv_role\":\"kv_both\",\"kv_connector_module_path\":\"tpu_inference.distributed.offload.tpu_offload_connector\"}' --port 8000 --enable-chunked-prefill --tensor-parallel-size 8 --seed 42 --enable_prefix_caching --gpu-memory-utilization 0.9"
+        - "vllm serve meta-llama/Llama-3.3-70B-Instruct --kv-transfer-config '{\"kv_connector\":\"TPUOffloadConnector\",\"kv_role\":\"kv_both\",\"kv_connector_module_path\":\"tpu_inference.offload.tpu_offload_connector\"}' --port 8000 --enable-chunked-prefill --tensor-parallel-size 8 --seed 42 --enable_prefix_caching --gpu-memory-utilization 0.9"
         env:
         - name: HUGGING_FACE_HUB_TOKEN
           valueFrom:
diff --git a/examples/offload/gke/pod_tpu_commons_cpu_offload.yaml b/examples/offload/gke/pod_tpu_commons_cpu_offload.yaml
@@ -18,7 +18,7 @@ spec:
     - --tensor_parallel_size=8
     - --max_model_len=1024
     - --kv-transfer-config
-    - '{"kv_connector":"TPUOffloadConnector","kv_connector_module_path":"tpu_inference.distributed.offload.tpu_offload_connector","kv_role":"kv_both"}'
+    - '{"kv_connector":"TPUOffloadConnector","kv_connector_module_path":"tpu_inference.offload.tpu_offload_connector","kv_role":"kv_both"}'
     env:
     - name: HUGGING_FACE_HUB_TOKEN
       valueFrom:
diff --git a/examples/offload/gke/pod_tpu_commons_cpu_offload_verification.yaml b/examples/offload/gke/pod_tpu_commons_cpu_offload_verification.yaml
@@ -25,7 +25,7 @@ spec:
     - --max_model_len=1024
     - --seed=42
     - --kv-transfer-config
-    - '{"kv_connector":"TPUOffloadConnector","kv_connector_module_path":"tpu_inference.distributed.offload.tpu_offload_connector","kv_role":"kv_both"}'
+    - '{"kv_connector":"TPUOffloadConnector","kv_connector_module_path":"tpu_inference.offload.tpu_offload_connector","kv_role":"kv_both"}'
     env:
     - name: HUGGING_FACE_HUB_TOKEN
       valueFrom:
diff --git a/examples/offload/gke/pod_tpu_host_offload_unit_tests.yaml b/examples/offload/gke/pod_tpu_host_offload_unit_tests.yaml
@@ -17,7 +17,7 @@ spec:
     command:
     - /bin/bash
     - -c
-    - "pytest -sv tests/distributed/offload/"
+    - "pytest -sv tests/offload/"
     env:
     - name: HUGGING_FACE_HUB_TOKEN
       valueFrom:
diff --git a/tests/offload/tpu_offload_accuracy_test.py b/tests/offload/tpu_offload_accuracy_test.py
@@ -40,8 +40,7 @@ def kv_transfer_config():
     return KVTransferConfig(
         kv_connector="TPUOffloadConnector",
         kv_role="kv_both",
-        kv_connector_module_path=
-        "tpu_inference.distributed.offload.tpu_offload_connector",
+        kv_connector_module_path="tpu_inference.offload.tpu_offload_connector",
     )
 
 
diff --git a/tests/offload/tpu_offload_connector_scheduler_test.py b/tests/offload/tpu_offload_connector_scheduler_test.py
@@ -9,7 +9,7 @@
 from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput
 from vllm.v1.request import Request
 
-from tpu_inference.distributed.offload.tpu_offload_connector import (
+from tpu_inference.offload.tpu_offload_connector import (
     RequestTracker, TPUOffloadConnectorScheduler)
 
 _DEFAULT_BLOCK_SIZE = 16
diff --git a/tests/offload/tpu_offload_connector_worker_test.py b/tests/offload/tpu_offload_connector_worker_test.py
@@ -15,13 +15,12 @@
 from jax.sharding import Mesh, NamedSharding, PartitionSpec
 from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole
 
-from tpu_inference.distributed.offload.tpu_offload_connector import (LoadSpec,
-                                                                     SaveSpec)
-from tpu_inference.distributed.offload.tpu_offload_connector import \
+from tpu_inference.logger import init_logger
+from tpu_inference.offload.tpu_offload_connector import LoadSpec, SaveSpec
+from tpu_inference.offload.tpu_offload_connector import \
     TPUOffloadConnector as CPUOffloadingConnector
-from tpu_inference.distributed.offload.tpu_offload_connector import (
+from tpu_inference.offload.tpu_offload_connector import (
     TPUOffloadConnectorMetadata, TPUReqMeta)
-from tpu_inference.logger import init_logger
 from tpu_inference.runner.tpu_runner import TPUModelRunner
 
 logger = init_logger(__name__)
diff --git a/tests/offload/tpu_offload_cpu_backend_test.py b/tests/offload/tpu_offload_cpu_backend_test.py
@@ -4,8 +4,8 @@
 
 import pytest
 
-from tpu_inference.distributed.offload.cpu_backend import LocalCPUBackend
-from tpu_inference.distributed.offload.utils import CpuChunkId
+from tpu_inference.offload.cpu_backend import LocalCPUBackend
+from tpu_inference.offload.utils import CpuChunkId
 
 
 # Helper to create a mock jax array with a specific size in bytes
diff --git a/tests/offload/tpu_offload_manager_test.py b/tests/offload/tpu_offload_manager_test.py
@@ -1,10 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
 import pytest
 
-from tpu_inference.distributed.offload.offload_manager import (
-    CPUChunkPool, LRUCacheManager, StagingBufferManager)
-from tpu_inference.distributed.offload.utils import ReqId
 from tpu_inference.logger import init_logger
+from tpu_inference.offload.offload_manager import (CPUChunkPool,
+                                                   LRUCacheManager,
+                                                   StagingBufferManager)
+from tpu_inference.offload.utils import ReqId
 
 logger = init_logger(__name__)
 
diff --git a/tests/offload/tpu_offload_utils_test.py b/tests/offload/tpu_offload_utils_test.py
@@ -7,8 +7,8 @@
 import numpy as np
 from jax.sharding import NamedSharding, PartitionSpec
 
-from tpu_inference.distributed.offload.utils import (
-    get_kv_cache_swap_fn, jitted_insert_kv_cache_slices)
+from tpu_inference.offload.utils import (get_kv_cache_swap_fn,
+                                         jitted_insert_kv_cache_slices)
 
 
 class TestTPUOffloadUtilsFn(unittest.TestCase):
diff --git a/tpu_inference/offload/__init__.py b/tpu_inference/offload/__init__.py
diff --git a/tpu_inference/offload/cpu_backend.py b/tpu_inference/offload/cpu_backend.py
@@ -1,12 +1,11 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import sys
 from collections import OrderedDict
 from typing import Any, Optional
 
-from tpu_inference.distributed.offload.utils import CpuChunkId
 from tpu_inference.logger import init_logger
+from tpu_inference.offload.utils import CpuChunkId
 
 logger = init_logger(__name__)
 
diff --git a/tpu_inference/offload/offload_manager.py b/tpu_inference/offload/offload_manager.py
@@ -1,14 +1,13 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 from collections import OrderedDict
 from dataclasses import dataclass
 from typing import Literal, Optional, Tuple
 
 from vllm.v1.core.kv_cache_utils import BlockHash
 
-from tpu_inference.distributed.offload.utils import CpuChunkId, ReqId
 from tpu_inference.logger import init_logger
+from tpu_inference.offload.utils import CpuChunkId, ReqId
 
 logger = init_logger(__name__)
 
diff --git a/tpu_inference/offload/tpu_offload_connector.py b/tpu_inference/offload/tpu_offload_connector.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """
 Scheduler side execution:
 TPUOffloadConnectorScheduler manages the state of KV cache loading and saving for
@@ -112,13 +111,14 @@
     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,
-    get_kv_cache_swap_fn, jitted_insert_kv_cache_slices)
 from tpu_inference.logger import init_logger
+from tpu_inference.offload.cpu_backend import LocalCPUBackend
+from tpu_inference.offload.offload_manager import (LRUCacheManager,
+                                                   StagingBufferManager)
+from tpu_inference.offload.utils import (CPU_OFFLOADING_SWAP_OP_TYPE,
+                                         CpuChunkId, KVCacheSwapFn, ReqId,
+                                         get_kv_cache_swap_fn,
+                                         jitted_insert_kv_cache_slices)
 from tpu_inference.runner.kv_cache_manager import KVCacheManager
 from tpu_inference.runner.tpu_runner import TPUModelRunner
 
diff --git a/tpu_inference/offload/utils.py b/tpu_inference/offload/utils.py
@@ -1,5 +1,4 @@
 # SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the LMCache project
 
 import functools
 import hashlib
diff --git a/tpu_inference/runner/kv_cache_manager.py b/tpu_inference/runner/kv_cache_manager.py
@@ -19,9 +19,8 @@
 
 from tpu_inference import utils
 from tpu_inference import utils as common_utils
-from tpu_inference.distributed.offload.utils import \
-    get_kv_connector_cache_layout
 from tpu_inference.logger import init_logger
+from tpu_inference.offload.utils import get_kv_connector_cache_layout
 from tpu_inference.runner import utils as runner_utils
 from tpu_inference.runner.input_batch import CachedRequestState, InputBatch
 from tpu_inference.runner.kv_cache import create_kv_caches
diff --git a/tpu_inference/worker/tpu_worker.py b/tpu_inference/worker/tpu_worker.py
@@ -292,7 +292,8 @@ def determine_available_memory(self) -> int:
 
         if self.vllm_config.kv_transfer_config is not None:
             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_transfer_config.kv_connector == "TPUOffloadConnector" and \
+               kv_transfer_config.kv_connector_module_path == "tpu_inference.offload.tpu_offload_connector":
                 # If kv offloading is enabled, we need to account for the memory used by the KV transfer buffer.
                 staging_buffer_pages = envs.TPU_OFFLOAD_NUM_STAGING_BLOCKS
 

Original file line number	Diff line number	Diff line change
`@@ -40,8 +40,7 @@ def kv_transfer_config():`
`40`	`40`	`return KVTransferConfig(`
`41`	`41`	`kv_connector="TPUOffloadConnector",`
`42`	`42`	`kv_role="kv_both",`
`43`		`- kv_connector_module_path=`
`44`		`- "tpu_inference.distributed.offload.tpu_offload_connector",`
	`43`	`+ kv_connector_module_path="tpu_inference.offload.tpu_offload_connector",`
`45`	`44`	`)`
`46`	`45`
`47`	`46`