Add miscellaneous updates (#8)

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,