fix

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

fix

Signed-off-by: 0oshowero0 <o0shower0o@outlook.com>

Author: 0oshowero0

tests/test_tensor_utils.py

@@ -0,0 +1,196 @@
+# Copyright 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
+# Copyright 2025 The TransferQueue Team
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for transfer_queue.utils.tensor_utils."""
+
+import pytest
+import torch
+
+from transfer_queue.utils.tensor_utils import (
+    allocate_empty_tensors,
+    compute_stride,
+    get_nbytes,
+    merge_contiguous_memory,
+)
+
+
+class TestComputeStride:
+    """Tests for compute_stride."""
+
+    def test_3d(self):
+        assert compute_stride((2, 3, 4)) == (12, 4, 1)
+
+    def test_1d(self):
+        assert compute_stride((5,)) == (1,)
+
+    def test_scalar(self):
+        assert compute_stride(()) == ()
+
+    def test_2d(self):
+        assert compute_stride((3, 5)) == (5, 1)
+
+
+class TestGetNbytes:
+    """Tests for get_nbytes."""
+
+    def test_basic(self):
+        dtypes = [torch.float32, torch.int32]
+        shapes = [(2, 3), (4,)]
+        result = get_nbytes(dtypes, shapes)
+        assert result == [2 * 3 * 4, 4 * 4]  # float32=4, int32=4
+
+    def test_scalar(self):
+        dtypes = [torch.float64]
+        shapes = [()]
+        result = get_nbytes(dtypes, shapes)
+        assert result == [8]  # scalar = 1 element
+
+    def test_list_shape(self):
+        dtypes = [torch.float32]
+        shapes = [[]]  # list instead of tuple
+        result = get_nbytes(dtypes, shapes)
+        assert result == [4]
+
+    def test_mixed_dtypes(self):
+        dtypes = [torch.float16, torch.float32, torch.int64]
+        shapes = [(10,), (10,), (10,)]
+        result = get_nbytes(dtypes, shapes)
+        assert result == [10 * 2, 10 * 4, 10 * 8]
+
+
+class TestAllocateEmptyTensors:
+    """Tests for allocate_empty_tensors."""
+
+    def test_basic(self):
+        dtypes = [torch.float32, torch.float32, torch.int32]
+        shapes = [(2, 3), (4,), (5,)]
+        tensors, ptrs, region_ptrs, region_sizes = allocate_empty_tensors(dtypes, shapes)
+
+        assert len(tensors) == 3
+        assert len(ptrs) == 3
+        assert len(region_ptrs) == 2  # float32 group + int32 group
+        assert len(region_sizes) == 2
+
+        # Same dtype tensors share the same underlying storage
+        assert tensors[0].untyped_storage().data_ptr() == region_ptrs[0]
+        assert tensors[1].untyped_storage().data_ptr() == region_ptrs[0]
+        assert tensors[2].untyped_storage().data_ptr() == region_ptrs[1]
+
+        # Shapes are correct
+        assert list(tensors[0].shape) == [2, 3]
+        assert list(tensors[1].shape) == [4]
+        assert list(tensors[2].shape) == [5]
+
+    def test_scalar(self):
+        dtypes = [torch.float32, torch.int32]
+        shapes = [(), ()]
+        tensors, ptrs, region_ptrs, region_sizes = allocate_empty_tensors(dtypes, shapes)
+
+        assert len(tensors) == 2
+        assert tensors[0].numel() == 1
+        assert tensors[1].numel() == 1
+        assert len(region_ptrs) == 2
+
+    def test_empty(self):
+        result = allocate_empty_tensors([], [])
+        assert result == ([], [], [], [])
+
+    def test_regions_complex(self):
+        """Mixed dtypes and shapes: verify region counts, sizes, and per-tensor offsets."""
+        dtypes = [
+            torch.float32,  # group 0: (2, 3) -> 6 elements
+            torch.int32,  # group 1: (4,) -> 4 elements
+            torch.float32,  # group 0: scalar -> 1 element
+            torch.float64,  # group 2: (2, 2) -> 4 elements
+            torch.int32,  # group 1: (3, 2) -> 6 elements
+        ]
+        shapes = [(2, 3), (4,), (), (2, 2), (3, 2)]
+        tensors, ptrs, region_ptrs, region_sizes = allocate_empty_tensors(dtypes, shapes)
+
+        # 3 dtype groups in insertion order: float32, int32, float64
+        assert len(region_ptrs) == 3
+        assert len(region_sizes) == 3
+        assert len(set(region_ptrs)) == 3  # distinct allocations
+
+        # float32 region: 6 + 1 = 7 elements * 4 bytes = 28 bytes
+        assert region_sizes[0] == 7 * 4
+        # int32 region: 4 + 6 = 10 elements * 4 bytes = 40 bytes
+        assert region_sizes[1] == 10 * 4
+        # float64 region: 4 elements * 8 bytes = 32 bytes
+        assert region_sizes[2] == 4 * 8
+
+        # Per-tensor ptrs must lie inside their respective regions
+        # tensor 0 (float32, shape (2,3), offset 0)
+        assert ptrs[0] == region_ptrs[0]
+        # tensor 1 (int32, shape (4,), offset 0)
+        assert ptrs[1] == region_ptrs[1]
+        # tensor 2 (float32, scalar, offset 6)
+        assert ptrs[2] == region_ptrs[0] + 6 * 4
+        # tensor 3 (float64, shape (2,2), offset 0)
+        assert ptrs[3] == region_ptrs[2]
+        # tensor 4 (int32, shape (3,2), offset 4)
+        assert ptrs[4] == region_ptrs[1] + 4 * 4
+
+
+class TestMergeContiguousMemory:
+    """Tests for merge_contiguous_memory."""
+
+    def test_basic_merge(self):
+        ptrs = [0, 10, 30]
+        sizes = [10, 20, 10]
+        merged_ptrs, merged_sizes = merge_contiguous_memory(ptrs, sizes)
+        # 0+10=10 (contiguous with 10), 10+20=30 (contiguous with 30) -> all merge into [0]
+        assert merged_ptrs == [0]
+        assert merged_sizes == [40]
+
+    def test_no_contiguous(self):
+        ptrs = [0, 100, 200]
+        sizes = [50, 50, 50]
+        merged_ptrs, merged_sizes = merge_contiguous_memory(ptrs, sizes)
+        assert merged_ptrs == [0, 100, 200]
+        assert merged_sizes == [50, 50, 50]
+
+    def test_unsorted_input(self):
+        ptrs = [100, 0, 50]
+        sizes = [50, 50, 50]
+        merged_ptrs, merged_sizes = merge_contiguous_memory(ptrs, sizes)
+        # After sorting: 0, 50, 100; all contiguous -> merge into [0]
+        assert merged_ptrs == [0]
+        assert merged_sizes == [150]
+
+    def test_single_region(self):
+        ptrs = [10]
+        sizes = [100]
+        merged_ptrs, merged_sizes = merge_contiguous_memory(ptrs, sizes)
+        assert merged_ptrs == [10]
+        assert merged_sizes == [100]
+
+    def test_empty(self):
+        assert merge_contiguous_memory([], []) == ([], [])
+
+    def test_mismatched_lengths_both_empty_not_triggered(self):
+        # If one is empty and other is not, should raise ValueError
+        with pytest.raises(ValueError, match="ptrs and sizes must have the same length"):
+            merge_contiguous_memory([], [10])
+
+        with pytest.raises(ValueError, match="ptrs and sizes must have the same length"):
+            merge_contiguous_memory([0], [])
+
+    def test_three_continuous(self):
+        ptrs = [0, 10, 20]
+        sizes = [10, 10, 10]
+        merged_ptrs, merged_sizes = merge_contiguous_memory(ptrs, sizes)
+        assert merged_ptrs == [0]
+        assert merged_sizes == [30]

transfer_queue/storage/clients/mooncake_client.py

@@ -24,7 +24,7 @@
 
 from transfer_queue.storage.clients.base import TransferQueueStorageKVClient
 from transfer_queue.storage.clients.factory import StorageClientFactory
-from transfer_queue.utils.tensor_utils import allocate_empty_tensors, get_nbytes, merge_continues_memory
+from transfer_queue.utils.tensor_utils import allocate_empty_tensors, get_nbytes, merge_contiguous_memory
 
 logger = logging.getLogger(__name__)
 logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
@@ -144,8 +144,8 @@ def put(self, keys: list[str], values: list[Any]) -> None:
     def _put_tensors_thread_worker(self, batch_keys: list[str], batch_tensors: list[Tensor]):
         """Worker thread for putting batch of tensors to MooncakeStore."""
 
-        batch_ptrs, batch_sizes, contiguous_tensors = self._preprocess_tensors_for_put(batch_tensors)
-        batch_ptr_reduced, batch_sizes_reduced = merge_continues_memory(batch_ptrs, batch_sizes)
+        batch_ptrs, batch_sizes, _contiguous_tensors = self._preprocess_tensors_for_put(batch_tensors)
+        batch_ptr_reduced, batch_sizes_reduced = merge_contiguous_memory(batch_ptrs, batch_sizes)
         self._register_all_buffers(batch_ptr_reduced, batch_sizes_reduced)
 
         try:
@@ -154,19 +154,19 @@ def _put_tensors_thread_worker(self, batch_keys: list[str], batch_tensors: list[
                 failed_indices = [j for j, r in enumerate(results) if r != 0]
                 error_codes = [results[j] for j in failed_indices]
                 raise RuntimeError(
-                    f"batch_put_tensor failed for indices {failed_indices} with error codes: {error_codes}"
+                    f"batch_upsert_from failed for indices {failed_indices} with error codes: {error_codes}"
                 )
         finally:
             self._unregister_all_buffers(batch_ptr_reduced)
 
-    def _put_bytes_thread_worker(self, batch_keys: list[str], batch_values: list[bytes]):
+    def _put_bytes_thread_worker(self, batch_keys: list[str], batch_values: list[Any]):
         """Worker thread for putting batch of non-tensors to MooncakeStore."""
 
         batch_values = [pickle.dumps(v, protocol=pickle.HIGHEST_PROTOCOL) for v in batch_values]
 
         ret = self._store.upsert_batch(batch_keys, batch_values, self.replica_config)
         if ret != 0:
-            raise RuntimeError(f"put_batch failed with error code: {ret}")
+            raise RuntimeError(f"upsert_batch failed with error code: {ret}")
 
     def get(
         self,
@@ -232,9 +232,11 @@ def _get_tensors_thread_worker(
         self, batch_keys: list[str], batch_shapes: list[tuple], batch_dtypes: list[torch.dtype], indexes: list[int]
     ) -> tuple[list[Tensor], list[int]]:
         batch_nbytes = get_nbytes(batch_dtypes, batch_shapes)
-        batch_buffer_tensors, batch_buffer_ptrs = allocate_empty_tensors(batch_dtypes, batch_shapes)
+        batch_buffer_tensors, batch_buffer_ptrs, region_ptrs, region_sizes = allocate_empty_tensors(
+            batch_dtypes, batch_shapes
+        )
 
-        self._register_all_buffers(batch_buffer_ptrs, batch_nbytes)
+        self._register_all_buffers(region_ptrs, region_sizes)
         try:
             ret_codes = self._store.batch_get_into(batch_keys, batch_buffer_ptrs, batch_nbytes)
             if len(ret_codes) != len(batch_keys):
@@ -243,7 +245,7 @@ def _get_tensors_thread_worker(
                 if ret < 0:
                     raise RuntimeError(f"batch_get_into failed for key `{batch_keys[i]}` with error code: {ret}")
         finally:
-            self._unregister_all_buffers(batch_buffer_ptrs)
+            self._unregister_all_buffers(region_ptrs)
 
         return batch_buffer_tensors, indexes
 
@@ -283,6 +285,11 @@ def _preprocess_tensors_for_put(values: list[Tensor]) -> tuple[list[Any], list[A
         size_list = []
         tensor_list = []  # hold reference for the contiguous tensor
         for t in values:
+            # TODO: support gpu direct rdma and use different data paths.
+            #       For GPU, it's more reasonable to perform data copy since
+            #       The register overhead is much higher than CPU
+            if t.device.type == "cuda":
+                t = t.cpu()
             t = t.contiguous()
             tensor_list.append(t)
             ptr_list.append(t.data_ptr())

transfer_queue/utils/tensor_utils.py

@@ -25,7 +25,9 @@
 logger.setLevel(os.getenv("TQ_LOGGING_LEVEL", logging.WARNING))
 
 
-def allocate_empty_tensors(dtypes: list[torch.dtype], shapes: list[tuple]) -> tuple[list[Tensor], list[int]]:
+def allocate_empty_tensors(
+    dtypes: list[torch.dtype], shapes: list[tuple]
+) -> tuple[list[Tensor], list[int], list[int], list[int]]:
     """Allocate empty tensors, grouping same dtypes into shared memory blocks.
 
     Instead of allocating each tensor separately, this function groups tensors
@@ -40,17 +42,19 @@ def allocate_empty_tensors(dtypes: list[torch.dtype], shapes: list[tuple]) -> tu
         A tuple containing:
             - List of tensors sharing memory within their dtype groups.
             - List of memory pointers (data_ptr) for each tensor.
+            - List of base pointers for each allocated memory region (one per dtype).
+            - List of total bytes for each allocated memory region (one per dtype).
 
     Example:
         >>> dtypes = [torch.float32, torch.float32, torch.int32, torch.float32]
         >>> shapes = [(10,), (20,), (5,), (15,)]
-        >>> tensors, ptrs = allocate_empty_tensors(dtypes, shapes)
+        >>> tensors, ptrs, region_ptrs, region_sizes = allocate_empty_tensors(dtypes, shapes)
         >>> # tensors[0], [1], [3] share the same dtype and memory block
     """
     assert len(dtypes) == len(shapes), "dtypes and shapes must have the same length"
 
     if len(dtypes) == 0:
-        return [], []
+        return [], [], [], []
 
     # Group indices by dtype
     dtype_groups: dict[torch.dtype, list[int]] = {}
@@ -61,6 +65,8 @@ def allocate_empty_tensors(dtypes: list[torch.dtype], shapes: list[tuple]) -> tu
 
     tensor_list = [torch.empty(()) for _ in range(len(dtypes))]
     ptr_list = [0] * len(dtypes)
+    region_ptrs: list[int] = []
+    region_sizes: list[int] = []
 
     # For each dtype group, allocate one big tensor and create views
     for dtype, indices in dtype_groups.items():
@@ -69,13 +75,15 @@ def allocate_empty_tensors(dtypes: list[torch.dtype], shapes: list[tuple]) -> tu
         shape_info = []  # Store (index, shape, num_elements, offset)
 
         for idx in indices:
-            shape = shapes[idx]
-            num_elements = reduce(operator.mul, shape)
+            shape = tuple(shapes[idx])
+            num_elements = reduce(operator.mul, shape, 1)
             shape_info.append((idx, shape, num_elements, total_elements))
             total_elements += num_elements
 
         # Allocate one big contiguous memory block for this dtype
         big_tensor = torch.empty(total_elements, dtype=dtype)
+        region_ptrs.append(big_tensor.data_ptr())
+        region_sizes.append(big_tensor.nbytes)
 
         # Create views into the big tensor for each small tensor
         for idx, shape, num_elements, offset in shape_info:
@@ -84,7 +92,7 @@ def allocate_empty_tensors(dtypes: list[torch.dtype], shapes: list[tuple]) -> tu
             tensor_list[idx] = small_tensor
             ptr_list[idx] = small_tensor.data_ptr()
 
-    return tensor_list, ptr_list
+    return tensor_list, ptr_list, region_ptrs, region_sizes
 
 
 def compute_stride(shape: tuple[int, ...]) -> tuple[int, ...]:
@@ -115,36 +123,37 @@ def get_nbytes(dtypes, shapes) -> list[int]:
     nbytes = []
     for i in range(len(dtypes)):
         elem_size = torch.tensor([], dtype=dtypes[i]).element_size()
-        numel = reduce(operator.mul, shapes[i])
+        shape = tuple(shapes[i])
+        numel = reduce(operator.mul, shape, 1)
         nbytes.append(elem_size * numel)
 
     return nbytes
 
 
-def merge_continues_memory(ptrs: list[int], sizes: list[int]) -> tuple[list[int], list[int]]:
-    """Merge continuous memory regions to reduce register_buffer overhead
+def merge_contiguous_memory(ptrs: list[int], sizes: list[int]) -> tuple[list[int], list[int]]:
+    """Merge contiguous memory regions to reduce register_buffer overhead
 
     Args:
         ptrs: List of memory pointers (starting addresses).
         sizes: List of memory region sizes corresponding to each pointer.
 
     Returns:
-        A tuple of (merged_ptrs, merged_sizes) where continuous regions
+        A tuple of (merged_ptrs, merged_sizes) where contiguous regions
         have been merged into single regions.
 
     Example:
-        >>> merge_continues_memory([0, 10, 30], [10, 20, 10])
+        >>> merge_contiguous_memory([0, 10, 30], [10, 20, 10])
         ([0, 30], [30, 10])
 
-        >>> merge_continues_memory([0, 5, 20], [5, 5, 10])
+        >>> merge_contiguous_memory([0, 5, 20], [5, 5, 10])
         ([0, 20], [10, 10])
     """
-    if not ptrs or not sizes:
-        return [], []
-
     if len(ptrs) != len(sizes):
         raise ValueError("ptrs and sizes must have the same length")
 
+    if not ptrs:
+        return [], []
+
     # Create list of (ptr, size) pairs and sort by pointer address
     regions = sorted(zip(ptrs, sizes, strict=False), key=lambda x: x[0])