biondizzle/vllm
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Use NCCL instead of ray for control-plane communication to remove serialization overhead (#2221)

Browse Source
This commit is contained in:
Zhuohan Li
2024-01-04 03:30:22 +08:00
committed by GitHub
parent 1066cbd152
commit fd4ea8ef5c
34 changed files with 524 additions and 262 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
tests/kernels/test_cache.py
View File

@@ -49,12 +49,13 @@ def test_copy_blocks(
src_blocks = random.sample(range(num_blocks), num_mappings)
remainig_blocks = list(set(range(num_blocks)) - set(src_blocks))
dst_blocks = random.sample(remainig_blocks, 2 * num_mappings)
block_mapping = {}
copy_src = []
copy_dst = []
for i in range(num_mappings):
src = src_blocks[i]
dst1 = dst_blocks[2 * i]
dst2 = dst_blocks[2 * i + 1]
block_mapping[src] = [dst1, dst2]
copy_src.append(src_blocks[i])
copy_dst.append(dst_blocks[2 * i])
copy_src.append(src_blocks[i])
copy_dst.append(dst_blocks[2 * i + 1])
# Create the KV caches.
key_caches, value_caches = kv_cache_factory(num_blocks, block_size,
@@ -66,15 +67,14 @@ def test_copy_blocks(
cloned_value_caches = [value_cache.clone() for value_cache in value_caches]
# Call the copy blocks kernel.
cache_ops.copy_blocks(key_caches, value_caches, block_mapping)
cache_ops.copy_blocks(key_caches, value_caches, copy_src, copy_dst)
# Run the reference implementation.
for src, dsts in block_mapping.items():
for dst in dsts:
for cloned_key_cache in cloned_key_caches:
cloned_key_cache[dst].copy_(cloned_key_cache[src])
for cloned_value_cache in cloned_value_caches:
cloned_value_cache[dst].copy_(cloned_value_cache[src])
for src, dst in zip(copy_src, copy_dst):
for cloned_key_cache in cloned_key_caches:
cloned_key_cache[dst].copy_(cloned_key_cache[src])
for cloned_value_cache in cloned_value_caches:
cloned_value_cache[dst].copy_(cloned_value_cache[src])
# Compare the results.
for key_cache, cloned_key_cache in zip(key_caches, cloned_key_caches):