Add nfsstore bandwidth testing script

Author: zhou-haitao

ucm/store/test/e2e/nfsstore_embed_fetch.py

@@ -49,13 +49,13 @@ def setup(
 
 
 def make_buffers(
-    block_number, device_id, batch_size, head_dim, block_len, block_layer, num_head
+    block_number, device_id, batch_size, head_dim, block_len, block_layer, num_head, kv
 ):
     hashes = [secrets.token_hex(16) for _ in range(block_number)]
     kv_caches = {}
     for i in range(block_layer):
         kv_caches[i] = torch.rand(
-            [1, block_number, block_len, num_head, head_dim],
+            [kv, block_number, block_len, num_head, head_dim],
             dtype=torch.bfloat16,
             device=f"cuda:{device_id}",
         )
@@ -73,10 +73,7 @@ def embed(
     store: UcmKVStoreBase,
     hashes: List[str],
     kvcaches: Dict[int, torch.Tensor],
-    num_tokens: int,
-    block_len: int,
-    block_layer: int,
-    block_dim: int,
+    mla: bool,
 ):
     start_time = time.perf_counter()
 
@@ -86,14 +83,23 @@ def embed(
     for i, hash_val in enumerate(hashes):
         offset = 0
         for layer_id, kv_layer in kvcaches.items():
-            tensor = kv_layer[0][i]  # kv=1
-            total_tensors.append(tensor)
+            k_tensor = kv_layer[0][i]  # kv=1
+            total_tensors.append(k_tensor)
             total_block_ids.append(hash_val)
             total_offsets.append(offset)
-            sz = tensor.numel() * tensor.element_size()
+            sz = k_tensor.numel() * k_tensor.element_size()
             offset += sz
             total_size += sz
 
+            if not mla:
+                v_tensor = kv_layer[1][i]
+                total_tensors.append(v_tensor)
+                total_block_ids.append(hash_val)
+                total_offsets.append(offset)
+                sz = v_tensor.numel() * v_tensor.element_size()
+                offset += sz
+                total_size += sz
+
     task = store.dump(total_block_ids, total_offsets, total_tensors)
     store.wait(task)
 
@@ -112,10 +118,7 @@ def fetch(
     store: UcmKVStoreBase,
     hashes: List[str],
     kvcaches: Dict[int, torch.Tensor],
-    num_tokens: int,
-    block_len: int,
-    block_layer: int,
-    block_dim: int,
+    mla: bool,
 ):
     start_time = time.perf_counter()
 
@@ -129,14 +132,23 @@ def fetch(
     for i, hash_val in enumerate(hashes):
         offset = 0
         for layer_id, kv_layer in kvcaches.items():
-            tensor = kv_layer[0][i]  # kv=1
+            k_tensor = kv_layer[0][i]  # kv=1
             block_ids.append(hash_val)
             offsets.append(offset)
-            tensors.append(tensor)
-            sz = tensor.numel() * tensor.element_size()
+            tensors.append(k_tensor)
+            sz = k_tensor.numel() * k_tensor.element_size()
             offset += sz
             total_size += sz
 
+            if not mla:
+                v_tensor = kv_layer[1][i]
+                block_ids.append(hash_val)
+                offsets.append(offset)
+                tensors.append(v_tensor)
+                sz = v_tensor.numel() * v_tensor.element_size()
+                offset += sz
+                total_size += sz
+
     task = store.load(block_ids, offsets, tensors)
     ret = store.wait(task)
     assert ret == 0, "Load operation failed"
@@ -163,6 +175,8 @@ def run(
     block_layer: int,
     head_size: int,
     block_elem_size: int,
+    kv: int,
+    mla: bool,
 ) -> Tuple[float, float, float, float, float, float]:
     """
     Run a single test with given parameters and return performance metrics.
@@ -195,6 +209,7 @@ def run(
             block_len,
             block_layer,
             num_head,
+            kv,
         )
 
         results = store.create(hashes[:batch_size])
@@ -204,10 +219,7 @@ def run(
             store,
             hashes[:batch_size],
             kvcaches,
-            num_tokens,
-            block_len,
-            block_layer,
-            block_dim,
+            mla,
         )
         store.commit(hashes[:batch_size], True)
 
@@ -217,10 +229,7 @@ def run(
             store,
             hashes[:batch_size],
             kvcaches,
-            num_tokens,
-            block_len,
-            block_layer,
-            block_dim,
+            mla,
         )
 
         w_bw_list.append(w_bw)

ucm/store/test/e2e/nfsstore_embed_fetch_run.py

@@ -64,7 +64,7 @@ def main():
         num_head_list = [1, 2, 4, 8]
 
     SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
-    csv_file = os.path.join(SCRIPT_DIR, "embed_fetch_result.csv")
+    csv_file = os.path.join(SCRIPT_DIR, "embed_fetch_result_all_6.csv")
     need_header = not os.path.exists(csv_file)
 
     with open(csv_file, "a", newline="", encoding="utf-8") as csv_fp:
@@ -124,6 +124,8 @@ def main():
                                 block_layer,
                                 head_size,
                                 block_elem_size,
+                                kv,
+                                mla,
                             ),
                         )