Revert "swap impl draft"

This reverts commit d43dc76.

Author: ZehaoLu98

vllm/distributed/kv_transfer/kv_connector/v1/swap_connector.py

@@ -19,7 +19,7 @@
 import torch
 
 from vllm import _custom_ops as ops
-from vllm.model_executor.layers.attention import Attention
+from vllm.attention.layer import Attention
 from vllm.config import VllmConfig, get_layers_from_vllm_config
 from vllm.distributed.kv_events import BlockRemoved, BlockStored, KVCacheEvent
 from vllm.distributed.kv_transfer.kv_connector.utils import yield_req_data
@@ -511,9 +511,6 @@ def __init__(self, spec: OffloadingSpec):
         self._gpu_tensors: dict[str, torch.Tensor] = {}
         self._cpu_tensors: dict[str, torch.Tensor] = {}
         self._kv_dim_before_num_blocks: dict[str, bool] = {}
-        self._block_size_bytes: dict[str, int] = {}
-        # True when all layers share a single GPU tensor
-        self._single_tensor_mode: bool = False
 
         # CUDA streams for per-layer transfers
         self._load_stream: torch.cuda.Stream | None = None
@@ -530,16 +527,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         Create per-layer CPU tensors and register handlers for bulk transfers.
         """
         layer_names = list(kv_caches.keys())
-
-        # Detect single-tensor mode: all layers share the same GPU tensor
-        gpu_ptrs = {t.data_ptr() for t in kv_caches.values()}
-        self._single_tensor_mode = (len(gpu_ptrs) == 1 and len(layer_names) > 1)
-        if self._single_tensor_mode:
-            logger.info(
-                "Single-tensor swap mode: %d layers share 1 GPU KV buffer",
-                len(layer_names),
-            )
-
         layers = get_layers_from_vllm_config(
             self.spec.vllm_config, Attention, layer_names
         )
@@ -548,15 +535,11 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
             for layer_name in layer_names
         }
 
-        # Register handlers for bulk transfers (prefix cache loads/stores).
-        # Skip in single-tensor mode: bulk handlers assume per-layer GPU
-        # tensors and won't work with a shared buffer. The per-layer
-        # load/store cycle handles all transfers instead.
-        if not self._single_tensor_mode:
-            for src_cls, dst_cls, handler in self.spec.get_handlers(
-                kv_caches, attn_backends
-            ):
-                self.worker.register_handler(src_cls, dst_cls, handler)
+        # Register handlers for bulk transfers (prefix cache loads/stores)
+        for src_cls, dst_cls, handler in self.spec.get_handlers(
+            kv_caches, attn_backends
+        ):
+            self.worker.register_handler(src_cls, dst_cls, handler)
 
         pin_memory = is_pin_memory_available()
         num_cpu_blocks = self.spec.num_blocks
@@ -597,24 +580,12 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
                 layer_name,
                 cpu_shape,
             )
-            cpu_tensor = torch.zeros(
+            self._cpu_tensors[layer_name] = torch.zeros(
                 cpu_shape,
                 dtype=gpu_tensor.dtype,
                 device="cpu",
                 pin_memory=pin_memory,
             )
-            self._cpu_tensors[layer_name] = cpu_tensor
-
-            # Compute block size in bytes for ops.swap_blocks (new 4-arg API).
-            # When kv_dim=True, shape is (2, num_blocks, ...) and we swap
-            # each K/V half separately via tensor[0]/tensor[1].
-            if self._kv_dim_before_num_blocks[layer_name]:
-                ref = cpu_tensor[0]
-            else:
-                ref = cpu_tensor
-            self._block_size_bytes[layer_name] = (
-                ref.element_size() * ref.stride(0)
-            )
 
         # Summary log after all CPU tensors are registered
         total_cpu_bytes = sum(
@@ -640,9 +611,6 @@ def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):
         )
 
     def handle_preemptions(self, preempted_req_ids: set[str]):
-        if self._single_tensor_mode:
-            return
-
         for job_id, transfer_spec in self._unsubmitted_store_jobs:
             success = self.worker.transfer_async(job_id, transfer_spec)
             assert success
@@ -655,11 +623,6 @@ def handle_preemptions(self, preempted_req_ids: set[str]):
 
     def start_kv_transfers(self, metadata: SwapConnectorMetadata):
         """Submit deferred stores and start prefix cache loads."""
-        # In single-tensor mode, skip bulk transfers entirely.
-        # Per-layer load/store handles everything.
-        if self._single_tensor_mode:
-            return
-
         # Submit deferred store jobs from the previous step
         for job_id, transfer_spec in self._unsubmitted_store_jobs:
             success = self.worker.transfer_async(job_id, transfer_spec)
@@ -740,14 +703,13 @@ def load_layer_from_cpu(
         cpu_tensor = self._cpu_tensors[layer_name]
         gpu_tensor = self._gpu_tensors[layer_name]
         kv_dim = self._kv_dim_before_num_blocks[layer_name]
-        block_size_bytes = self._block_size_bytes[layer_name]
 
         with torch.cuda.stream(self._load_stream):
             if kv_dim:
-                ops.swap_blocks(cpu_tensor[0], gpu_tensor[0], block_size_bytes, src_to_dst_tensor)
-                ops.swap_blocks(cpu_tensor[1], gpu_tensor[1], block_size_bytes, src_to_dst_tensor)
+                ops.swap_blocks(cpu_tensor[0], gpu_tensor[0], src_to_dst_tensor)
+                ops.swap_blocks(cpu_tensor[1], gpu_tensor[1], src_to_dst_tensor)
             else:
-                ops.swap_blocks(cpu_tensor, gpu_tensor, block_size_bytes, src_to_dst_tensor)
+                ops.swap_blocks(cpu_tensor, gpu_tensor, src_to_dst_tensor)
 
         # Must synchronize: attention needs the data to be ready
         self._load_stream.synchronize()
@@ -814,14 +776,13 @@ def store_layer_to_cpu(
         cpu_tensor = self._cpu_tensors[layer_name]
         gpu_tensor = self._gpu_tensors[layer_name]
         kv_dim = self._kv_dim_before_num_blocks[layer_name]
-        block_size_bytes = self._block_size_bytes[layer_name]
 
         with torch.cuda.stream(self._store_stream):
             if kv_dim:
-                ops.swap_blocks(gpu_tensor[0], cpu_tensor[0], block_size_bytes, src_to_dst_tensor)
-                ops.swap_blocks(gpu_tensor[1], cpu_tensor[1], block_size_bytes, src_to_dst_tensor)
+                ops.swap_blocks(gpu_tensor[0], cpu_tensor[0], src_to_dst_tensor)
+                ops.swap_blocks(gpu_tensor[1], cpu_tensor[1], src_to_dst_tensor)
             else:
-                ops.swap_blocks(gpu_tensor, cpu_tensor, block_size_bytes, src_to_dst_tensor)
+                ops.swap_blocks(gpu_tensor, cpu_tensor, src_to_dst_tensor)
             # Record event for the load stream to wait on
             if self._store_event is None:
                 self._store_event = torch.Event()
@@ -836,9 +797,6 @@ def wait_for_all_stores(self):
 
     def prepare_store_kv(self, metadata: SwapConnectorMetadata):
         """Prepare bulk store jobs for the scheduler's reqs_to_store."""
-        if self._single_tensor_mode:
-            return
-
         for req_id, transfer_spec in metadata.reqs_to_store.items():
             job_id = self._generate_job_id()
             self._jobs[job_id] = (req_id, True)
@@ -848,15 +806,10 @@ def prepare_store_kv(self, metadata: SwapConnectorMetadata):
     def get_finished(
         self, finished_req_ids: set[str]
     ) -> tuple[set[str], set[str]]:
-        # In single-tensor mode, no bulk jobs are submitted
-        if self._single_tensor_mode:
-            return set(), set()
-
         finished_sending = set()
         finished_recving = set()
-        for result in self.worker.get_finished():
-            assert result.success
-            job_id = result.job_id
+        for job_id, success in self.worker.get_finished():
+            assert success
             req_id, store = self._jobs.pop(job_id)
             if store:
                 req_jobs = self._store_jobs[req_id]

vllm/v1/core/kv_cache_utils.py

@@ -1145,95 +1145,12 @@ def get_kv_cache_config_from_groups(
                 KVCacheTensor(size=page_size * num_blocks, shared_by=shared_by)
             )
 
-    config = KVCacheConfig(
+    return KVCacheConfig(
         num_blocks=num_blocks,
         kv_cache_tensors=kv_cache_tensors,
         kv_cache_groups=kv_cache_groups,
     )
 
-    # In single-GPU-tensor swap mode, all layers share 1 GPU KV cache tensor.
-    # num_blocks is determined by CPU capacity, not GPU.
-    config = _maybe_apply_single_gpu_tensor(vllm_config, config,
-                                            available_memory)
-
-    return config
-
-
-def _maybe_apply_single_gpu_tensor(
-    vllm_config: VllmConfig,
-    config: KVCacheConfig,
-    available_memory: int,
-) -> KVCacheConfig:
-    """
-    If single_gpu_tensor swap mode is enabled, replace the N per-layer
-    KVCacheTensors with a single shared tensor. num_blocks is capped by
-    both CPU capacity and available GPU memory.
-    """
-    kv_transfer_config = vllm_config.kv_transfer_config
-    if kv_transfer_config is None:
-        return config
-    extra = kv_transfer_config.kv_connector_extra_config or {}
-    if not extra.get("single_gpu_tensor", False):
-        return config
-
-    cpu_bytes_to_use = extra.get("cpu_bytes_to_use")
-    if not cpu_bytes_to_use:
-        raise ValueError(
-            "cpu_bytes_to_use must be set when single_gpu_tensor is enabled"
-        )
-    cpu_bytes_to_use = int(cpu_bytes_to_use)
-
-    # Collect all layer names and compute page size
-    all_layer_names: list[str] = []
-    for group in config.kv_cache_groups:
-        all_layer_names.extend(group.layer_names)
-    num_layers = len(all_layer_names)
-    assert num_layers > 0
-
-    # All tensors should have the same page size
-    page_sizes = set()
-    for group in config.kv_cache_groups:
-        page_sizes.add(group.kv_cache_spec.page_size_bytes)
-    assert len(page_sizes) == 1, (
-        "single_gpu_tensor mode requires uniform page sizes across groups"
-    )
-    page_size = page_sizes.pop()
-
-    # num_blocks determined by CPU capacity, capped by GPU memory.
-    # CPU holds all layers' data; GPU holds just 1 layer's worth.
-    cpu_num_blocks = int(cpu_bytes_to_use // (page_size * num_layers))
-    gpu_num_blocks = int(available_memory // page_size)
-    num_blocks = min(cpu_num_blocks, gpu_num_blocks)
-    assert num_blocks > 0, (
-        f"Cannot allocate any KV blocks. "
-        f"cpu_bytes_to_use={cpu_bytes_to_use} -> {cpu_num_blocks} blocks, "
-        f"available_gpu_memory={available_memory} -> {gpu_num_blocks} blocks, "
-        f"num_layers={num_layers}, page_size={page_size}"
-    )
-
-    bottleneck = "GPU" if gpu_num_blocks < cpu_num_blocks else "CPU"
-    logger.info(
-        "Single-GPU-tensor swap mode: %d blocks (limited by %s), "
-        "%d layers sharing 1 GPU tensor of %s, CPU backing %s",
-        num_blocks,
-        bottleneck,
-        num_layers,
-        format_gib(page_size * num_blocks),
-        format_gib(cpu_bytes_to_use),
-    )
-
-    # One GPU tensor shared by ALL layers
-    single_tensor = KVCacheTensor(
-        size=page_size * num_blocks,
-        shared_by=all_layer_names,
-    )
-
-    return KVCacheConfig(
-        num_blocks=num_blocks,
-        kv_cache_tensors=[single_tensor],
-        kv_cache_groups=config.kv_cache_groups,
-    )
-
 
 def unify_hybrid_kv_cache_specs(kv_cache_spec: dict[str, KVCacheSpec]):
     """

vllm/v1/core/sched/scheduler.py

@@ -939,7 +939,7 @@ def schedule(self) -> SchedulerOutput:
         )
 
         with record_function_or_nullcontext(nvtx_label):
-            self._update_after_schedule(scheduler_output)
+                self._update_after_schedule(scheduler_output)
         return scheduler_output
 
     def _preempt_request(self, request: Request, timestamp: float) -> None:

vllm/v1/kv_offload/cpu_swap.py

@@ -52,16 +52,9 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
         }
         assert len(page_sizes) == 1
         page_size_bytes = page_sizes.pop()
-
-        # Use actual layer count from groups (not len(kv_cache_tensors),
-        # which is 1 in single-GPU-tensor mode).
-        num_layers = sum(
-            len(g.layer_names)
-            for g in kv_cache_config.kv_cache_groups
-        )
         kv_bytes_per_block = (
             page_size_bytes
-            * num_layers
+            * len(kv_cache_config.kv_cache_tensors)
             * vllm_config.parallel_config.world_size
         )
         kv_bytes_per_offloaded_block = kv_bytes_per_block * (
@@ -74,9 +67,8 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
             else 0
         )
 
-        # In single-GPU-tensor mode, all layers share one GPU tensor,
-        # so total_gpu_blocks is just the num_blocks from the config.
-        # In normal mode, sum across tensors.
+        # Calculate the total number of GPU blocks that could be allocated.
+        # The CPU must be able to hold all of them.
         total_gpu_blocks = sum(
             kv_cache_tensor.size // page_size_bytes
             for kv_cache_tensor in kv_cache_config.kv_cache_tensors
@@ -87,9 +79,6 @@ def __init__(self, vllm_config: VllmConfig, kv_cache_config: KVCacheConfig):
             total_gpu_blocks + block_size_factor - 1
         ) // block_size_factor
 
-        # In single-GPU-tensor mode, the GPU tensor is small (1 layer)
-        # and the CPU holds all layers' data. The assertion is naturally
-        # satisfied since num_blocks was computed from cpu_bytes_to_use.
         assert self.num_blocks >= required_offloaded_blocks, (
             f"CPU swap mode requires enough CPU memory to hold all KV cache. "
             f"CPU can hold {self.num_blocks} offloaded blocks but "

vllm/v1/worker/gpu/attn_utils.py

@@ -167,6 +167,58 @@ def build_slot_mappings_by_layer(
     return slot_mappings_by_layer
 
 
+def init_kv_cache_with_offloading(
+    runner_kv_caches: list[torch.Tensor],
+    forward_context: dict[str, Any],
+    kv_cache_config: KVCacheConfig,
+    attn_backends: dict[str, AttentionBackend],
+    gpu_device: torch.device,
+    num_gpu_buffer_layers: int = 2,
+) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+    """Initialize KV cache with per-layer CPU offloading.
+
+    Returns:
+        A tuple of (gpu_kv_caches, cpu_kv_caches), each mapping
+        layer_name -> reshaped KV cache tensor.
+    """
+    gpu_raw, cpu_raw = _allocate_kv_cache_with_offloading(
+        kv_cache_config, gpu_device, num_gpu_buffer_layers,
+    )
+    gpu_kv_caches = _reshape_kv_cache(kv_cache_config, gpu_raw, attn_backends)
+    cpu_kv_caches = _reshape_kv_cache(kv_cache_config, cpu_raw, attn_backends)
+
+    # Bind the GPU buffer tensors so attention layers can use them
+    bind_kv_cache(gpu_kv_caches, forward_context, runner_kv_caches)
+
+    return gpu_kv_caches, cpu_kv_caches
+
+
+def init_kv_cache_with_offloading(
+    runner_kv_caches: list[torch.Tensor],
+    forward_context: dict[str, Any],
+    kv_cache_config: KVCacheConfig,
+    attn_backends: dict[str, AttentionBackend],
+    gpu_device: torch.device,
+    num_gpu_buffer_layers: int = 2,
+) -> tuple[dict[str, torch.Tensor], dict[str, torch.Tensor]]:
+    """Initialize KV cache with per-layer CPU offloading.
+
+    Returns:
+        A tuple of (gpu_kv_caches, cpu_kv_caches), each mapping
+        layer_name -> reshaped KV cache tensor.
+    """
+    gpu_raw, cpu_raw = _allocate_kv_cache_with_offloading(
+        kv_cache_config, gpu_device, num_gpu_buffer_layers,
+    )
+    gpu_kv_caches = _reshape_kv_cache(kv_cache_config, gpu_raw, attn_backends)
+    cpu_kv_caches = _reshape_kv_cache(kv_cache_config, cpu_raw, attn_backends)
+
+    # Bind the GPU buffer tensors so attention layers can use them
+    bind_kv_cache(gpu_kv_caches, forward_context, runner_kv_caches)
+
+    return gpu_kv_caches, cpu_kv_caches
+
+
 def build_attn_metadata(
     attn_groups: list[list[AttentionGroup]],
     num_reqs: int,

vllm_profile/vllm_profile.py

@@ -58,7 +58,7 @@ def main():
         profiler_config = None
 
     ktc = KVTransferConfig(
-        kv_connector="SwapConnector",
+        kv_connector="LMCacheConnectorV1",
         kv_role="kv_both",
         kv_connector_extra_config={
             "cpu_bytes_to_use": 64 * 1024 ** 3,  # 64 GiB of pinned CPU RAM for KV cache