Add chunkSize

Author: qyh111

ucm/integration/vllm/uc_connector.py

@@ -113,6 +113,8 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
         self.num_head = vllm_config.model_config.get_num_kv_heads(
             vllm_config.parallel_config
         )
+        self.chunk_size = 256
+        self.blocks_per_chunk = self.chunk_size // self.block_size
         self.head_size = vllm_config.model_config.get_head_size()
         if (
             self._vllm_config.kv_transfer_config is not None
@@ -139,8 +141,9 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
                 config_base
                 * self.num_layers
                 * (1 if self.is_mla else self.num_head * self.total_tp_size * 2)
-            )
-            config["io_size"] = config_base * (1 if self.is_mla else self.num_head)
+            ) * self.blocks_per_chunk
+            self.io_size = config_base * (1 if self.is_mla else self.num_head) * self.blocks_per_chunk
+            config["io_size"] = self.io_size
             logger.info(
                 "kv_block_size = %d, io_size = %d,",
                 config["kv_block_size"],
@@ -160,8 +163,6 @@ def __init__(self, vllm_config: "VllmConfig", role: KVConnectorRole):
                     "use_layerwise"
                 ]
             )
-        self.chunk_size = 256
-        self.blocks_per_chunk = self.chunk_size // self.block_size
 
     def _init_kv_caches_from_forward_context(self, forward_context: "ForwardContext"):
         for layer_name in forward_context.no_compile_layers:
@@ -207,25 +208,54 @@ def DataOffset(self, kv_layer, rank, layer_id, is_v):
                 layer_size * layer_id + layer_size / self.total_tp_size * self.rank
             )
 
-    def get_tensor_and_offset_layerwise(
+    def get_pointers_and_offset_layerwise(
         self, vllm_block_ids_tensors: List[torch.Tensor], kv_layer: torch.Tensor, layer_name: str
+    ) -> tuple[List[List[int]], List[int]]:
+        k_pointer_lists = []
+        k_offsets = []
+        v_pointer_lists = []
+        v_offsets = []
+        layer_id = self._extract_layer_index(layer_name)
+
+        for vllm_block_ids_tensor in vllm_block_ids_tensors:
+            vllm_block_ids = vllm_block_ids_tensor.tolist()
+            k_pointer_list = []
+            k_data_offset = self.DataOffset(kv_layer, self.rank, layer_id, False)
+            for vllm_block_id in vllm_block_ids:
+                if self.is_mla:
+                    k_pointer_list.append(kv_layer[vllm_block_id].data_ptr())
+                else:
+                    k_pointer_list.append(kv_layer[0][vllm_block_id].data_ptr())
+            k_pointer_lists.append(k_pointer_list)
+            k_offsets.append(k_data_offset)
+            if not self.is_mla:
+                v_pointer_list = []
+                v_data_offset = self.DataOffset(kv_layer, self.rank, layer_id, True)
+                for vllm_block_id in vllm_block_ids:
+                    v_pointer_list.append(kv_layer[1][vllm_block_id].data_ptr())
+                v_offsets.append(v_data_offset)
+                v_pointer_lists.append(v_pointer_list)
+        return k_pointer_lists + v_pointer_lists, k_offsets + v_offsets
+
+    def get_tensor_and_offset_layerwise(
+        self, vllm_block_ids: List[int], kv_layer: torch.Tensor, layer_name: str
     ) -> tuple[List[torch.Tensor], List[int]]:
         k_tensors = []
         k_offsets = []
         v_tensors = []
         v_offsets = []
         layer_id = self._extract_layer_index(layer_name)
 
-        for vllm_block_ids_tensor in vllm_block_ids_tensors:
+        for blk_id in vllm_block_ids:
             k_data_offset = self.DataOffset(kv_layer, self.rank, layer_id, False)
             if self.is_mla:
-                k_tensors.append(kv_layer[vllm_block_ids_tensor])
+                k_tensors.append(kv_layer[blk_id])
             else:
-                k_tensors.append(kv_layer[0][vllm_block_ids_tensor])
+                k_tensors.append(kv_layer[0][blk_id])
             k_offsets.append(k_data_offset)
             if not self.is_mla:
                 v_data_offset = self.DataOffset(kv_layer, self.rank, layer_id, True)
-                v_tensors.append(kv_layer[1][vllm_block_ids_tensor])
+                v_tensors.append(kv_layer[1][blk_id])
                 v_offsets.append(v_data_offset)
         return k_tensors + v_tensors, k_offsets + v_offsets
 
@@ -277,18 +307,20 @@ def start_load_kv(self, forward_context: "ForwardContext", **kwargs) -> None:
                     vllm_block_ids_tensor.tolist()
                 )
             for layer_name, kv_layer in self.kv_caches.items():
-                tensors, offsets = self.get_tensor_and_offset_layerwise(
+                pointers_list, offsets = self.get_pointers_and_offset_layerwise(
                     vllm_block_ids_tensors, kv_layer, layer_name
                 )
-                k_task_id = self.connector.load(
-                    storage_block_ids, offsets[:blocks_len], tensors[:blocks_len]
+                size = [self.io_size] * blocks_len
+                k_task_id = self.connector.fetch_data(
+                    storage_block_ids, offsets[:blocks_len], pointers_list[:blocks_len], size
                 )
                 v_task_id = None
                 if not self.is_mla:
-                    v_task_id = self.connector.load(
+                    v_task_id = self.connector.fetch_data(
                         storage_block_ids,
                         offsets[blocks_len:],
-                        tensors[blocks_len:],
+                        pointers_list[blocks_len:],
+                        size
                     )
                 if request.request_id not in self.layerwise_load_tasks:
                     self.layerwise_load_tasks[request.request_id] = {}
@@ -404,7 +436,7 @@ def save_kv_layer(
             storage_block_ids = [block[0] for block in request.dump_blocks]
             vllm_block_ids_tensors = [block[1] for block in request.dump_blocks]  # [5, 8, 12]
             blocks_len = len(storage_block_ids)
-            tensors, offsets = self.get_tensor_and_offset_layerwise(
+            pointers_list, offsets = self.get_pointers_and_offset_layerwise(
                 vllm_block_ids_tensors, kv_layer, layer_name
             )
 
@@ -413,18 +445,18 @@ def save_kv_layer(
             elif kv_layer[0].device.type == "cuda":
                 torch.cuda.current_stream().synchronize()
 
-            for block_id, offset, tensor in zip(
-                storage_block_ids, offsets[:blocks_len], tensors[:blocks_len]
+            for block_id, offset, pointers in zip(
+                storage_block_ids, offsets[:blocks_len], pointers_list[:blocks_len]
             ):
-                task = self.connector.dump([block_id], [offset], [tensor])
+                task = self.connector.dump_data([block_id], [offset], [pointers], [self.io_size])
                 self.dump_tasks.setdefault(request.request_id, {}).setdefault(
                     block_id, []
                 ).append(task)
             if not self.is_mla:
-                for block_id, offset, tensor in zip(
-                    storage_block_ids, offsets[blocks_len:], tensors[blocks_len:]
+                for block_id, offset, pointer_lists in zip(
+                    storage_block_ids, offsets[blocks_len:], pointers_list[blocks_len:]
                 ):
-                    task = self.connector.dump([block_id], [offset], [tensor])
+                    task = self.connector.dump_data([block_id], [offset], [pointer_lists], [self.io_size])
                     self.dump_tasks.setdefault(request.request_id, {}).setdefault(
                         block_id, []
                     ).append(task)
@@ -465,23 +497,23 @@ def wait_for_tasks():
             vllm_block_ids_tensors = [block[1] for block in request.dump_blocks]
             blocks_len = len(storage_block_ids)
             for layer_name, kv_layer in self.kv_caches.items():
-                tensors, offsets = self.get_tensor_and_offset_layerwise(
+                pointers_list, offsets = self.get_pointers_and_offset_layerwise(
                     vllm_block_ids_tensors, kv_layer, layer_name
                 )
-                for block_id, offset, tensor in zip(
-                    storage_block_ids, offsets[:blocks_len], tensors[:blocks_len]
+                for block_id, offset, pointers in zip(
+                    storage_block_ids, offsets[:blocks_len], pointers_list[:blocks_len]
                 ):
-                    task = self.connector.dump([block_id], [offset], [tensor])
+                    task = self.connector.dump_data([block_id], [offset], [pointers], [self.io_size])
                     self.dump_tasks.setdefault(request.request_id, {}).setdefault(
                         block_id, []
                     ).append(task)
                 if not self.is_mla:
-                    for block_id, offset, tensor in zip(
+                    for block_id, offset, pointers in zip(
                         storage_block_ids,
                         offsets[blocks_len:],
-                        tensors[blocks_len:],
+                        pointers_list[blocks_len:],
                     ):
-                        task = self.connector.dump([block_id], [offset], [tensor])
+                        task = self.connector.dump_data([block_id], [offset], [pointers], [self.io_size])
                         self.dump_tasks.setdefault(request.request_id, {}).setdefault(
                             block_id, []
                         ).append(task)
@@ -633,7 +665,7 @@ def hash_request_tokens(
                 start_position=start_position,
             )
             self._need_load_reqs[request.request_id] = []
-            return num_lookup_hits * self.block_size, True
+            return num_lookup_hits * self.block_size - num_computed_tokens, True
 
         # When all the tokens are cached in ssd or hbm,
         # we need to recompute the last token. This if condition will be removed
@@ -650,7 +682,7 @@ def hash_request_tokens(
             start_position=start_position,
         )
 
-        return num_lookup_hits * self.block_size, False
+        return num_lookup_hits * self.block_size - num_computed_tokens, False
 
     def update_state_after_alloc(
         self, request: "Request", blocks: "KVCacheBlocks", num_external_tokens: int

ucm/store/dramstore/cpy/dramstore.py.cc

@@ -78,8 +78,12 @@ class DRAMStorePy : public DRAMStore {
         auto length = lengths.begin();
         while ((blockId != blockIds.end()) && (offset != offsets.end()) &&
                (address != addresses.end()) && (length != lengths.end())) {
+            std::vector<uintptr_t> addr_vec;
+            for (auto addr_item : address->cast<py::list>()) {
+                addr_vec.push_back(addr_item.cast<uintptr_t>());
+            }
             task.Append(blockId->cast<std::string>(), offset->cast<size_t>(),
-                        address->cast<uintptr_t>(), length->cast<size_t>());
+                        std::move(addr_vec), length->cast<size_t>());
             blockId++;
             offset++;
             address++;

ucm/store/localstore/cpy/localstore.py.cc

@@ -78,8 +78,12 @@ class LocalStorePy : public LocalStore {
         auto length = lengths.begin();
         while ((blockId != blockIds.end()) && (offset != offsets.end()) &&
                (address != addresses.end()) && (length != lengths.end())) {
+            std::vector<uintptr_t> addr_vec;
+            for (auto addr_item : address->cast<py::list>()) {
+                addr_vec.push_back(addr_item.cast<uintptr_t>());
+            }
             task.Append(blockId->cast<std::string>(), offset->cast<size_t>(),
-                        address->cast<uintptr_t>(), length->cast<size_t>());
+                        std::move(addr_vec), length->cast<size_t>());
             blockId++;
             offset++;
             address++;

ucm/store/nfsstore/cc/domain/trans/posix_queue.cc

@@ -103,7 +103,13 @@ Status PosixQueue::D2S(Task::Shard& shard, const Device& device)
         return Status::OutOfMemory();
     }
     auto hub = shard.buffer.get();
-    auto status = device->D2HSync((std::byte*)hub, (std::byte*)shard.address, shard.length);
+    auto dAddr = new std::byte*[shard.address.size()];
+    auto hAddr = new std::byte*[shard.address.size()];
+    for (size_t i = 0; i < shard.address.size(); i++) {
+        hAddr[i] = (std::byte*)hub + i * shard.length / shard.address.size();
+        dAddr[i] = (std::byte*)shard.address[i];
+    }
+    auto status = device->D2HBatchSync(hAddr, const_cast<const std::byte**>(dAddr), shard.address.size(), shard.length / shard.address.size());
     if (status.Failure()) { return status; }
     auto path = this->layout_->DataFilePath(shard.block, true);
     return File::Write(path, shard.offset, shard.length, (uintptr_t)hub);
@@ -120,21 +126,27 @@ Status PosixQueue::S2D(Task::Shard& shard, const Device& device)
     auto path = this->layout_->DataFilePath(shard.block, false);
     auto status = File::Read(path, shard.offset, shard.length, (uintptr_t)hub);
     if (status.Failure()) { return status; }
-    return device->H2DAsync((std::byte*)shard.address, (std::byte*)hub, shard.length);
+    auto dAddr = new std::byte*[shard.address.size()];
+    auto hAddr = new std::byte*[shard.address.size()];
+    for (size_t i = 0; i < shard.address.size(); i++) {
+        hAddr[i] = (std::byte*)hub + i * shard.length / shard.address.size();
+        dAddr[i] = (std::byte*)shard.address[i];
+    }
+    return device->H2DBatchSync(dAddr, const_cast<const std::byte**>(hAddr), shard.address.size(), shard.length / shard.address.size());
 }
 
 Status PosixQueue::H2S(Task::Shard& shard)
 {
     auto path = this->layout_->DataFilePath(shard.block, true);
-    auto aligned = IsAligned(shard.offset) && IsAligned(shard.length) && IsAligned(shard.address);
-    return File::Write(path, shard.offset, shard.length, shard.address, aligned);
+    auto aligned = IsAligned(shard.offset) && IsAligned(shard.length) && IsAligned(shard.address[0]);
+    return File::Write(path, shard.offset, shard.length, shard.address[0], aligned);
 }
 
 Status PosixQueue::S2H(Task::Shard& shard)
 {
     auto path = this->layout_->DataFilePath(shard.block, false);
-    auto aligned = IsAligned(shard.offset) && IsAligned(shard.length) && IsAligned(shard.address);
-    return File::Read(path, shard.offset, shard.length, shard.address, aligned);
+    auto aligned = IsAligned(shard.offset) && IsAligned(shard.length) && IsAligned(shard.address[0]);
+    return File::Read(path, shard.offset, shard.length, shard.address[0], aligned);
 }
 
 } // namespace UC

ucm/store/nfsstore/cpy/nfsstore.py.cc

@@ -91,8 +91,12 @@ class NFSStorePy : public NFSStore {
         auto length = lengths.begin();
         while ((blockId != blockIds.end()) && (offset != offsets.end()) &&
                (address != addresses.end()) && (length != lengths.end())) {
+            std::vector<uintptr_t> addr_vec;
+            for (auto addr_item : address->cast<py::list>()) {
+                addr_vec.push_back(addr_item.cast<uintptr_t>());
+            }
             task.Append(blockId->cast<std::string>(), offset->cast<size_t>(),
-                        address->cast<uintptr_t>(), length->cast<size_t>());
+                        std::move(addr_vec), length->cast<size_t>());
             blockId++;
             offset++;
             address++;
@@ -123,8 +127,6 @@ PYBIND11_MODULE(ucmnfsstore, module)
     config.def_readwrite("transferBufferNumber", &UC::NFSStorePy::Config::transferBufferNumber);
     config.def_readwrite("transferTimeoutMs", &UC::NFSStorePy::Config::transferTimeoutMs);
     config.def_readwrite("tempDumpDirEnable", &UC::NFSStorePy::Config::tempDumpDirEnable);
-    config.def_readwrite("hotnessEnable", &UC::NFSStorePy::Config::hotnessEnable);
-    config.def_readwrite("hotnessInterval", &UC::NFSStorePy::Config::hotnessInterval);
     store.def(py::init<>());
     store.def("CCStoreImpl", &UC::NFSStorePy::CCStoreImpl);
     store.def("Setup", &UC::NFSStorePy::Setup);

ucm/store/nfsstore/nfsstore_connector.py

@@ -51,6 +51,7 @@ def __init__(self, config: Dict):
         if transfer_enable:
             param.transferDeviceId = config["device"]
             param.transferIoSize = config["io_size"]
+            param.transferStreamNumber = config.get("transfer_stream_number", 128)
         ret = self.store.Setup(param)
         if ret != 0:
             msg = f"Failed to initialize ucmnfsstore, errcode: {ret}."

ucm/store/task/task_shard.h

@@ -47,13 +47,13 @@ class Task {
         Location location;
         std::string block;
         size_t offset;
-        uintptr_t address;
+        std::vector<uintptr_t> address;
         size_t length;
         size_t owner;
         std::shared_ptr<void> buffer;
         std::function<void(void)> done;
         Shard(const Type type, const Location location, const std::string& block,
-              const size_t offset, const uintptr_t address, const size_t length, const size_t owner)
+              const size_t offset, const std::vector<uintptr_t> address, const size_t length, const size_t owner)
             : type{type}, location{location}, block{block}, offset{offset}, address{address},
               length{length}, owner{owner}, buffer{nullptr}, done{nullptr}
         {
@@ -86,7 +86,7 @@ class Task {
     auto Id() const noexcept { return id_; }
     auto StartTp() const noexcept { return startTp_; }
     auto Str() const noexcept { return fmt::format("{},{},{},{}", id_, brief_, number_, size_); }
-    void Append(const std::string& block, const size_t offset, const uintptr_t address,
+    void Append(const std::string& block, const size_t offset, const std::vector<uintptr_t> address,
                 const size_t length)
     {
         shards_.emplace_back(type_, location_, block, offset, address, length, id_);

ucm/store/ucmstore.py

@@ -129,7 +129,7 @@ def fetch_data(
         self,
         block_ids: List[str],
         offset: List[int],
-        dst_addr: List[int],
+        dst_addr: List[List[int]],
         size: List[int],
     ) -> Task:
         """
@@ -150,7 +150,7 @@ def dump_data(
         self,
         block_ids: List[str],
         offset: List[int],
-        src_addr: List[int],
+        src_addr: List[List[int]],
         size: List[int],
     ) -> Task:
         """