Use NCCL instead of ray for control-plane communication to remove serialization overhead

csrc/cache.h

@@ -8,12 +8,14 @@
 void swap_blocks(
   torch::Tensor& src,
   torch::Tensor& dst,
-  const std::map<int64_t, int64_t>& block_mapping);
+  const std::vector<int64_t>& src_block_numbers,
+  const std::vector<int64_t>& dst_block_numbers);
 
 void copy_blocks(
   std::vector<torch::Tensor>& key_caches,
   std::vector<torch::Tensor>& value_caches,
-  const std::map<int64_t, std::vector<int64_t>>& block_mapping);
+  const std::vector<int64_t>& src_block_numbers,
+  const std::vector<int64_t>& dst_block_numbers);
 
 void reshape_and_cache(
   torch::Tensor& key,

csrc/cache_kernels.cu

@@ -12,7 +12,9 @@
 void swap_blocks(
   torch::Tensor& src,
   torch::Tensor& dst,
-  const std::map<int64_t, int64_t>& block_mapping) {
+  const std::vector<int64_t>& src_block_numbers,
+  const std::vector<int64_t>& dst_block_numbers) {
+  assert(src_block_numbers.size() == dst_block_numbers.size());
   torch::Device src_device = src.device();
   torch::Device dst_device = dst.device();
   cudaMemcpyKind memcpy_type;
@@ -35,9 +37,9 @@ void swap_blocks(
   const int64_t block_size_in_bytes = src.element_size() * src[0].numel();
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
   // NOTE(woosuk): This can be slow if the number of blocks is large.
-  for (const auto& pair : block_mapping) {
-    int64_t src_block_number = pair.first;
-    int64_t dst_block_number = pair.second;
+  for (int64_t i = 0; i < src_block_numbers.size(); ++i) {
+    int64_t src_block_number = src_block_numbers[i];
+    int64_t dst_block_number = dst_block_numbers[i];
     int64_t src_offset = src_block_number * block_size_in_bytes;
     int64_t dst_offset = dst_block_number * block_size_in_bytes;
     cudaMemcpyAsync(
@@ -85,7 +87,8 @@ __global__ void copy_blocks_kernel(
 void copy_blocks(
   std::vector<torch::Tensor>& key_caches,
   std::vector<torch::Tensor>& value_caches,
-  const std::map<int64_t, std::vector<int64_t>>& block_mapping) {
+  const std::vector<int64_t>& src_block_numbers,
+  const std::vector<int64_t>& dst_block_numbers) {
   int num_layers = key_caches.size();
   TORCH_CHECK(num_layers == value_caches.size());
   if (num_layers == 0) {
@@ -104,12 +107,10 @@ void copy_blocks(
   }
   // Create block mapping array.
   std::vector<int64_t> block_mapping_vec;
-  for (const auto& pair : block_mapping) {
-    int64_t src_block_number = pair.first;
-    for (int64_t dst_block_number : pair.second) {
-      block_mapping_vec.push_back(src_block_number);
-      block_mapping_vec.push_back(dst_block_number);
-    }
+  assert(src_block_numbers.size() == dst_block_numbers.size());
+  for (int i = 0; i < src_block_numbers.size(); ++i) {
+    block_mapping_vec.push_back(src_block_numbers[i]);
+    block_mapping_vec.push_back(dst_block_numbers[i]);
   }
   int64_t* block_mapping_array = block_mapping_vec.data();
   int num_pairs = block_mapping_vec.size() / 2;
@@ -252,12 +253,12 @@ __global__ void gather_cached_kv_kernel(
     for (int i = threadIdx.x; i < num_tokens; i += blockDim.x) {
       const int tgt_key_idx = token_idx * key_stride + i;
       const int tgt_value_idx = token_idx * value_stride + i;
-  
+
       const int head_idx = i / head_size;
       const int head_offset = i % head_size;
       const int x_idx = head_offset / x;  // the offset of the [head_size/x] dimension
       const int x_offset = head_offset % x;
-  
+
       const int src_key_idx = block_idx * num_heads * (head_size / x) * block_size * x
                               + head_idx * (head_size / x) * block_size * x
                               + x_idx * block_size * x

docs/source/models/adding_model.rst

@@ -59,7 +59,7 @@ Next, you need to rewrite the :code:`forward` methods of your model by following
     +    kv_caches: List[KVCache],
     +    input_metadata: InputMetadata,
     +    cache_events: Optional[List[torch.cuda.Event]],
-    +) -> SamplerOutput:
+    +) -> Optional[SamplerOutput]:
 
 3. Update the code by considering that :code:`input_ids` and :code:`positions` are now flattened tensors.
 4. Replace the attention operation with either :code:`PagedAttention`, :code:`PagedAttentionWithRoPE`, or :code:`PagedAttentionWithALiBi` depending on the model's architecture.

examples/offline_inference_long.py

@@ -0,0 +1,17 @@
+from vllm import LLM, SamplingParams
+
+# Sample prompts.
+prompts = [
+    "hi" * 90000,
+]
+# Create a sampling params object.
+sampling_params = SamplingParams(temperature=0.0)
+
+# Create an LLM.
+llm = LLM(model="mistralai/Mistral-7B-v0.1", max_model_len=160000)
+# Generate texts from the prompts. The output is a list of RequestOutput objects
+# that contain the prompt, generated text, and other information.
+outputs = llm.generate(prompts, sampling_params)
+# Print the outputs.
+for output in outputs:
+    print(f"Prompt len: {len(output.prompt_token_ids)}, Generated text: {output.outputs[0].text!r}")

playground/test_ray_placement_group.py

@@ -0,0 +1,65 @@
+import time
+import os
+
+# Import placement group APIs.
+from ray.util.placement_group import (
+    placement_group,
+    placement_group_table,
+    remove_placement_group,
+)
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+# Initialize Ray.
+import ray
+
+class NormalActor:
+    def __init__(self, index):
+        self.index = index
+        pass
+
+    def log_message(self):
+        import torch
+        print("NormalActor", self.index, os.getpid(), torch.cuda.is_available(), ray.get_gpu_ids())
+
+class AllocationActor:
+    def __init__(self, pg):
+        self.placement_group = pg
+        self.a2 = ray.remote(num_cpus=1)(NormalActor).options(
+            scheduling_strategy=PlacementGroupSchedulingStrategy(
+                placement_group=pg,
+                placement_group_bundle_index=1,
+            )
+        ).remote(1)
+        self.a3 = ray.remote(num_gpus=1, num_cpus=0)(NormalActor).options(
+            scheduling_strategy=PlacementGroupSchedulingStrategy(
+                placement_group=pg,
+                placement_group_bundle_index=2,
+            )
+        ).remote(2)
+
+    def log_message(self):
+        print("AllocationActor", os.getpid())
+        ray.get([self.a2.log_message.remote(), self.a3.log_message.remote()])
+
+
+def main():
+    # Create a single node Ray cluster with 2 CPUs and 2 GPUs.
+    ray.init(num_cpus=2, num_gpus=1)
+
+    print(ray.cluster_resources())
+
+    # Reserve a placement group of 1 bundle that reserves 1 CPU and 1 GPU.
+    pg = placement_group([{"CPU": 1}, {"CPU": 1}, {"GPU": 1, "CPU": 0, "node:__internal_head__": 1e-2}])
+
+    ray.get(pg.ready())
+    a1 = ray.remote(num_cpus=1)(AllocationActor).options(
+        scheduling_strategy=PlacementGroupSchedulingStrategy(
+            placement_group=pg,
+            placement_group_bundle_index=0,
+        )
+    ).remote(pg)
+
+    ray.get(a1.log_message.remote())
+    print(ray.available_resources())
+
+main()