fix test case

Signed-off-by: Cruz Zhao <[email protected]>

Author: zxpdemonio

scripts/test_tensor_api.py

@@ -292,7 +292,7 @@ def test_04_put_get_into(self):
         """Verify basic put and get into functionality."""
         key = "get_into_test"
         tensor = torch.randn(1024, 1024, dtype=torch.float32)
-        buffer_spacing = 64 * 1024 * 1024  # 1GB per tensor slot
+        buffer_spacing = 64 * 1024 * 1024
         total_buffer_size = buffer_spacing
 
         # Allocate large contiguous buffer
@@ -322,7 +322,7 @@ def test_05_batch_put_get_into(self):
         """Zero copy Batch Get."""
         num_tensors = 4
         keys, tensors = generate_tensors(num_tensors, 8)
-        buffer_spacing = 1024 * 1024 * 1024  # 1GB per tensor slot
+        buffer_spacing = 64 * 1024 * 1024  # 1GB per tensor slot
         batch_size = len(keys)
         total_buffer_size = buffer_spacing * batch_size
 
@@ -378,7 +378,7 @@ def test_06_put_get_into_with_tp(self):
         registered_buffers = []  # Keep track of (ptr, size) for cleanup
 
         for rank in range(tp_size):
-            buffer_spacing = 1024 * 1024 * 1024  # 1GB per tensor slot
+            buffer_spacing = 64 * 1024 * 1024  # 1GB per tensor slot
             total_buffer_size = buffer_spacing
 
             # Allocate buffer for this rank
@@ -445,7 +445,7 @@ def test_07_batch_put_get_into_with_tp(self):
 
         for rank in range(tp_size):
             batch_size = len(keys)
-            buffer_spacing = 1024 * 1024 * 1024  # 1GB per tensor slot
+            buffer_spacing = 64 * 1024 * 1024  # 1GB per tensor slot
             total_buffer_size = buffer_spacing * batch_size
 
             # Allocate buffer for this rank
@@ -585,7 +585,8 @@ def test_benchmark_02_tp_batch(self):
 
     def test_benchmark_03_batch_put_get_into(self):
         """Benchmark: Zero copy Batch Get."""
-        buffer_spacing = 512 * 1024 * 1024  # 1GB per tensor slot
+        self.store.remove_all()
+        buffer_spacing = 300 * 1024 * 1024  # 1GB per tensor slot
         batch_size = len(self.keys)
         total_buffer_size = buffer_spacing * batch_size
 
@@ -645,7 +646,8 @@ def test_benchmark_04_batch_put_get_into_with_tp(self):
         tp_size = 4
         split_dim = 0
         batch_size = len(self.keys)
-        buffer_spacing = 512 * 1024 * 1024  # 1GB per tensor slot
+        self.store.remove_all()
+        buffer_spacing = 300 * 1024 * 1024  # 1GB per tensor slot
 
         # Allocate and register a separate buffer for each TP rank
         rank_buffers = []  # Store metadata for cleanup