Merge pull request #4 from robertgshaw2-redhat/rob-changes

Rob changes

Author: robertgshaw2-redhat

examples/offline_inference/disaggrated-prefill-v1/decode_example.py

@@ -17,8 +17,8 @@
 sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
 
 llm = LLM(
-    model="meta-llama/Llama-3.1-8B-Instruct",
-    enforce_eager=True,
+    model="meta-llama/Llama-3.2-1B-Instruct",
+    enforce_eager=False,
     gpu_memory_utilization=0.8,
     kv_transfer_config=KVTransferConfig.from_cli(
         '{"kv_connector":"SharedStorageConnector","kv_role":"kv_both",'

examples/offline_inference/disaggrated-prefill-v1/prefill_example.py

@@ -14,8 +14,8 @@
 
 sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
 
-llm = LLM(model="meta-llama/Llama-3.1-8B-Instruct",
-          enforce_eager=True,
+llm = LLM(model="meta-llama/Llama-3.2-1B-Instruct",
+          enforce_eager=False,
           gpu_memory_utilization=0.8,
           kv_transfer_config=KVTransferConfig.from_cli(
               '{"kv_connector":"SharedStorageConnector","kv_role":"kv_both", '

examples/offline_inference/disaggrated-prefill-v1/run.sh

@@ -1,5 +1,5 @@
 rm -rf local_storage/
 rm output.txt
 
-VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 prefill_example.py
-VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=0 python3 decode_example.py
+VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=5 python3 prefill_example.py
+VLLM_ENABLE_V1_MULTIPROCESSING=0 CUDA_VISIBLE_DEVICES=5 python3 decode_example.py

vllm/attention/layer.py

@@ -181,11 +181,6 @@ def forward(
         context using
         `vllm.forward_context.get_forward_context().attn_metadata`.
         """
-
-        # KVConnector: start async saving kvs to connector
-        # to the layers KV cache before running attention.
-        wait_for_kv_layer_from_connector(self.layer_name)
-
         if self.calculate_kv_scales:
             attn_metadata = get_forward_context().attn_metadata
             if attn_metadata.enable_kv_scales_calculation:
@@ -236,10 +231,6 @@ def forward(
                 output = torch.ops.vllm.unified_attention(
                     query, key, value, self.layer_name)
 
-        # KVConnector: start saving kvs to the connector.
-        # NOTE: forward_context completion will block until
-        # this operation is completed.
-        maybe_save_kv_layer_to_connector(self.layer_name, self.kv_cache)
         return output
 
     def calc_kv_scales(self, query, key, value):
@@ -358,7 +349,7 @@ def wait_for_kv_layer_from_connector(layer_name: str):
 
 def maybe_save_kv_layer_to_connector(
     layer_name: str,
-    kv_cache: List[torch.Tensor],
+    kv_cache_layer: List[torch.Tensor],
 ):
     if not has_kv_transfer_group() or not is_v1_kv_transfer_group():
         return
@@ -370,7 +361,6 @@ def maybe_save_kv_layer_to_connector(
     if attn_metadata is None:
         return
 
-    kv_cache_layer = kv_cache[forward_context.virtual_engine]
     connector.save_kv_layer(layer_name, kv_cache_layer, attn_metadata)
 
 
@@ -380,11 +370,17 @@ def unified_attention(
     value: torch.Tensor,
     layer_name: str,
 ) -> torch.Tensor:
+    wait_for_kv_layer_from_connector(layer_name)
+
     forward_context: ForwardContext = get_forward_context()
     attn_metadata = forward_context.attn_metadata
     self = forward_context.no_compile_layers[layer_name]
     kv_cache = self.kv_cache[forward_context.virtual_engine]
-    return self.impl.forward(self, query, key, value, kv_cache, attn_metadata)
+    output = self.impl.forward(self, query, key, value, kv_cache,
+                               attn_metadata)
+
+    maybe_save_kv_layer_to_connector(layer_name, kv_cache)
+    return output
 
 
 def unified_attention_fake(
@@ -412,6 +408,7 @@ def unified_attention_with_output(
     output: torch.Tensor,
     layer_name: str,
 ) -> None:
+    wait_for_kv_layer_from_connector(layer_name)
     forward_context: ForwardContext = get_forward_context()
     attn_metadata = forward_context.attn_metadata
     self = forward_context.no_compile_layers[layer_name]
@@ -424,6 +421,8 @@ def unified_attention_with_output(
                       attn_metadata,
                       output=output)
 
+    maybe_save_kv_layer_to_connector(layer_name, kv_cache)
+
 
 def unified_attention_with_output_fake(
     query: torch.Tensor,

vllm/distributed/kv_transfer/kv_connector/v1/shared_storage_connector.py

@@ -196,6 +196,9 @@ def extract_kv_from_layer(
 
             Assume the shape of the layer is (2, num_pages, page_size, xxx).
             """
+            # TODO(rob): make this compatible with MLA.
+
+            assert layer.shape[0] == 2
             num_pages, page_size = layer.shape[1], layer.shape[2]
             return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping,
                                                                ...]
@@ -208,7 +211,7 @@ def extract_kv_from_layer(
                     layer_name, request.token_ids)
                 kv_cache = extract_kv_from_layer(kv_layer,
                                                  request.slot_mapping)
-                tensors = {"kv_cache": kv_cache.cpu().detach()}
+                tensors = {"kv_cache": kv_cache.detach().cpu()}
                 safetensors.torch.save_file(tensors, filename)
 
     def wait_for_save(self):