Add memory analyzer & utomatically configure KV cache size (#6)

server.py

@@ -3,6 +3,7 @@ from typing import List
 
 from cacheflow.master.frontend import Frontend
 from cacheflow.master.scheduler import Scheduler
+from cacheflow.models import get_memory_analyzer
 from cacheflow.worker.controller import Controller
 
 parser = argparse.ArgumentParser(description='CacheFlow server')
@@ -10,17 +11,25 @@ parser.add_argument('--model', type=str, default='facebook/opt-125m', help='mode
 parser.add_argument('--num-nodes', type=int, default=1, help='number of nodes')
 parser.add_argument('--num-workers', type=int, default=1, help='number of workers per node')
 parser.add_argument('--block-size', type=int, default=8, choices=[8, 16], help='token block size')
-# TODO(woosuk): Add an analytical model to determine the maximum number of GPU/CPU blocks.
-parser.add_argument('--num-gpu-blocks', type=int, default=1024, help='number of GPU blocks (per GPU)')
-parser.add_argument('--num-cpu-blocks', type=int, default=32, help='number of CPU blocks (per GPU)')
 # NOTE(woosuk): If FlashAttention is used, the float data type is not supported.
 parser.add_argument('--dtype', type=str, default='half', choices=['half', 'float'], help='data type')
 # TODO(woosuk): Support fine-grained seeds (e.g., seed per request).
 parser.add_argument('--seed', type=int, default=0, help='random seed')
+parser.add_argument('--max-batch-size', type=int, default=2048, help='maximum number of batched tokens')
 args = parser.parse_args()
 
 
 def main():
+    memory_analyzer = get_memory_analyzer(
+        model_name=args.model,
+        block_size=args.block_size,
+        dtype=args.dtype,
+    )
+    num_gpu_blocks = memory_analyzer.get_max_num_gpu_blocks(
+        max_num_batched_tokens=args.max_batch_size)
+    num_cpu_blocks = memory_analyzer.get_max_num_cpu_blocks()
+    print(f'# GPU blocks: {num_gpu_blocks}, # CPU blocks: {num_cpu_blocks}')
+
     # Create a controller for each node.
     controllers: List[Controller] = []
     for i in range(args.num_nodes):
@@ -29,8 +38,8 @@ def main():
             num_workers=args.num_workers,
             model_name=args.model,
             block_size=args.block_size,
-            num_gpu_blocks=args.num_gpu_blocks,
-            num_cpu_blocks=args.num_cpu_blocks,
+            num_gpu_blocks=num_gpu_blocks,
+            num_cpu_blocks=num_cpu_blocks,
             dtype=args.dtype,
             seed=args.seed,
         )
@@ -47,8 +56,9 @@ def main():
         frontend=frontend,
         controllers=controllers,
         block_size=args.block_size,
-        num_gpu_blocks=args.num_gpu_blocks,
-        num_cpu_blocks=args.num_cpu_blocks,
+        num_gpu_blocks=num_gpu_blocks,
+        num_cpu_blocks=num_cpu_blocks,
+        max_num_batched_tokens=args.max_batch_size,
     )
     # Connect the controllers.
     for i in range(len(controllers) - 1):