Add miscellaneous updates (#8)

cacheflow/master/scheduler.py

@@ -158,8 +158,8 @@ class Scheduler:
         # 3. Join new sequences if possible.
         # NOTE: Here we implicitly assume FCFS scheduling.
         # TODO(woosuk): Add a batching policy to control the batch size.
+        self._fetch_inputs()
         if not self.swapped:
-            self._fetch_inputs()
             for i, seq_group in enumerate(self.pending):
                 num_prompt_tokens = seq_group.seqs[0].get_len()
                 if self.block_manager.can_allocate(seq_group):
@@ -211,12 +211,13 @@ class Scheduler:
             input_seq_groups.append(input_seq_group)
 
         # 5. Execute the first stage of the pipeline.
-        self.controllers[0].execute_stage(
-            input_seq_groups,
-            blocks_to_swap_in,
-            blocks_to_swap_out,
-            blocks_to_copy,
-        )
+        if (input_seq_groups or blocks_to_swap_in or blocks_to_swap_out):
+            self.controllers[0].execute_stage(
+                input_seq_groups,
+                blocks_to_swap_in,
+                blocks_to_swap_out,
+                blocks_to_copy,
+            )
 
     def post_step(
         self,

cacheflow/models/attention.py

@@ -12,7 +12,7 @@ from cacheflow.models import InputMetadata
 class OPTCacheFlowAttention(nn.Module):
 
     def __init__(self, scale: float) -> None:
-        super().__init__()
+        super(OPTCacheFlowAttention, self).__init__()
         self.scale = float(scale)
 
         self.flash_attn = FlashAttention(softmax_scale=self.scale)
@@ -106,8 +106,8 @@ class OPTCacheFlowAttention(nn.Module):
         output = output.view(-1, num_heads, head_size)
 
         # Compute the attention op for prompts.
-        if input_metadata.num_prompts > 0:
-            num_prompt_tokens = sum(input_metadata.prompt_lens)
+        num_prompt_tokens = input_metadata.num_prompt_tokens
+        if num_prompt_tokens > 0:
             self.multi_query_kv_attention(
                 output[:num_prompt_tokens],
                 query[:num_prompt_tokens],
@@ -126,10 +126,9 @@ class OPTCacheFlowAttention(nn.Module):
 
         if input_metadata.num_generation_tokens > 0:
             # Compute the attention op for generation tokens.
-            start_idx = sum(input_metadata.prompt_lens)
             self.single_query_cached_kv_attention(
-                output[start_idx:],
-                query[start_idx:],
+                output[num_prompt_tokens:],
+                query[num_prompt_tokens:],
                 key_cache,
                 value_cache,
                 input_metadata)

cacheflow/models/memory_analyzer.py

@@ -5,7 +5,7 @@ from cacheflow.models.utils import get_cpu_memory
 from cacheflow.models.utils import get_dtype_size
 from cacheflow.models.utils import get_gpu_memory
 
-_GiB = 1 << 30    
+_GiB = 1 << 30
 
 
 class CacheFlowMemoryAnalyzer:
@@ -117,9 +117,19 @@ class OPTMemoryAnalyzer(CacheFlowMemoryAnalyzer):
 
     def get_max_num_cpu_blocks(
         self,
-        memory_utilization: float = 0.25,
+        swap_space: int,
     ) -> int:
+        swap_space = swap_space * _GiB
         cpu_memory = get_cpu_memory()
-        usable_memory = int(memory_utilization * cpu_memory)
-        max_num_blocks = usable_memory // self._get_cache_block_size()
+        if swap_space > 0.8 * cpu_memory:
+            raise ValueError(f'The swap space ({swap_space / _GiB:.2f} GiB) '
+                             'takes more than 80% of the available memory '
+                             f'({cpu_memory / _GiB:.2f} GiB).'
+                             'Please check the swap space size.')
+        if swap_space > 0.5 * cpu_memory:
+            print(f'WARNING: The swap space ({swap_space / _GiB:.2f} GiB) '
+                  'takes more than 50% of the available memory '
+                  f'({cpu_memory / _GiB:.2f} GiB).'
+                  'This may slow the system performance.')
+        max_num_blocks = swap_space // self._get_cache_block_size()
         return max_num_blocks

cacheflow/models/sample.py

@@ -11,7 +11,7 @@ from cacheflow.sequence import SequenceOutputs
 class Sampler(nn.Module):
 
     def __init__(self) -> None:
-        super().__init__()
+        super(Sampler, self).__init__()
 
     def forward(
         self,

cacheflow/worker/worker.py

@@ -191,6 +191,13 @@ class Worker:
         else:
             cache_events = None
 
+        # If there is no input, we don't need to execute the model.
+        if not input_seq_groups:
+            if cache_events is not None:
+                for event in cache_events:
+                    event.wait()
+            return {}
+
         # Prepare input tensors.
         input_tokens, input_positions, input_metadata = self.prepare_inputs(
             input_seq_groups)