Support block size 32 (#35)

cacheflow/master/block_manager.py

@@ -15,9 +15,9 @@ class BlockAllocator:
         block_size: int,
         num_blocks: int,
     ) -> None:
-        if block_size not in [8, 16]:
+        if block_size not in [8, 16, 32]:
             raise ValueError(f'Unsupported block size: {block_size}'
-                             'The block size must be either 8 or 16.')
+                             'The block size must be one of {8, 16, 32}.')
         self.device = device
         self.block_size = block_size
         self.num_blocks = num_blocks

cacheflow/master/server.py

@@ -174,7 +174,7 @@ def add_server_arguments(parser: argparse.ArgumentParser):
     parser.add_argument('--pipeline-parallel-size', '-pp', type=int, default=1, help='number of pipeline stages')
     parser.add_argument('--tensor-parallel-size', '-tp', type=int, default=1, help='number of tensor parallel replicas')
     # KV cache arguments
-    parser.add_argument('--block-size', type=int, default=8, choices=[8, 16], help='token block size')
+    parser.add_argument('--block-size', type=int, default=8, choices=[8, 16, 32], help='token block size')
     # 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).