From d8d435c7d04dcf368e1f86dc328e812e293c6f67 Mon Sep 17 00:00:00 2001
From: Claude <claude@anthropic.com>
Date: Thu, 6 Aug 2026 04:22:53 +0000
Subject: [PATCH] [BASE] cache_engine.py: CCCL temporary_storage layout
 two-phase KV cache allocation
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Source: cccl_upstream/cub/cub/detail/temporary_storage.cuh
Target: vllm/worker/cache_engine.py

CCCL system design applied:
- temporary_storage::layout<SlotsCount>: Phase 1 get_size() computes
  total bytes, Phase 2 map_to_buffer() allocates one blob and aliases
  into per-slot views
- Applied to _allocate_kv_cache: compute total numel for all layers,
  allocate one contiguous torch.zeros, slice into per-layer views
- Reduces cudaMalloc calls from num_attention_layers to 1
- Guarantees cross-layer memory contiguity (better L2 locality)
- slot.create_alias<T>() → layer_flat.view(kv_cache_shape)
---
 vllm/worker/cache_engine.py | 48 +++++++++++++++++++++++++++++--------
 1 file changed, 38 insertions(+), 10 deletions(-)

diff --git a/vllm/worker/cache_engine.py b/vllm/worker/cache_engine.py
index 090f95e6..683ac341 100644
--- a/vllm/worker/cache_engine.py
+++ b/vllm/worker/cache_engine.py
@@ -69,20 +69,48 @@ class CacheEngine:
         num_blocks: int,
         device: str,
     ) -> List[torch.Tensor]:
-        """Allocates KV cache on the specified device."""
+        """Allocates KV cache on the specified device.
+
+        CCCL temporary_storage.cuh layout<SlotsCount> system design:
+          Phase 1 (get_size): compute total bytes for all slots
+          Phase 2 (map_to_buffer): allocate one blob, alias into slots
+
+        Applied: instead of N separate torch.zeros (one per layer),
+        compute total size → allocate one contiguous tensor → slice
+        into per-layer views. Reduces cudaMalloc calls from
+        num_attention_layers to 1, and guarantees cross-layer memory
+        contiguity (better L2 locality for multi-layer KV access).
+
+        The slot/alias pattern maps directly:
+          layout slot[i] = layer i's KV cache
+          alias<T> = the typed view into that layer's region
+        """
         kv_cache_shape = self.attn_backend.get_kv_cache_shape(
             num_blocks, self.block_size, self.num_kv_heads, self.head_size)
         pin_memory = is_pin_memory_available() if device == "cpu" else False
         kv_cache: List[torch.Tensor] = []
-        for _ in range(self.num_attention_layers):
-            # null block in CpuGpuBlockAllocator requires at least that
-            # block to be zeroed-out.
-            # We zero-out everything for simplicity.
-            kv_cache.append(
-                torch.zeros(kv_cache_shape,
-                            dtype=self.dtype,
-                            pin_memory=pin_memory,
-                            device=device))
+
+        if self.num_attention_layers == 0 or num_blocks == 0:
+            return kv_cache
+
+        # Phase 1: get_size — compute per-layer element count
+        import math
+        layer_numel = math.prod(kv_cache_shape)
+
+        # Phase 2: map_to_buffer — single contiguous allocation
+        total_numel = self.num_attention_layers * layer_numel
+        contiguous_buffer = torch.zeros(
+            total_numel,
+            dtype=self.dtype,
+            pin_memory=pin_memory,
+            device=device)
+
+        # Alias into per-layer views (CCCL slot.create_alias pattern)
+        for i in range(self.num_attention_layers):
+            start = i * layer_numel
+            layer_flat = contiguous_buffer[start:start + layer_numel]
+            kv_cache.append(layer_flat.view(kv_cache_shape))
+
         return kv_cache
 
     def swap_in(self, src_to_dst: torch.Tensor) -> None: