Support various block sizes

benchmark/benchmark_text_completion.py

@@ -268,6 +268,7 @@ def get_sampling_dir_name(
             f'{model_name}-tp{args.tensor_parallel_size}',
             sample_dir,
             'cacheflow',
+            f'block{args.block_size}',
             f'req-rate-{args.request_rate}',
             f'seed{args.seed}',
             f'duration-{args.duration}',

cacheflow/master/block_manager.py

@@ -15,9 +15,6 @@ def __init__(
         block_size: int,
         num_blocks: int,
     ) -> None:
-        if block_size not in [8, 16, 32]:
-            raise ValueError(f'Unsupported block size: {block_size}'
-                             '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/scheduler.py

@@ -125,7 +125,8 @@ def _schedule(
 
         # Swap in the sequence groups in the SWAPPED state if possible.
         self.swapped = self.policy.sort_by_priority(now, self.swapped)
-        while self.swapped:
+        # FCFS
+        while self.swapped and not blocks_to_swap_out:
             seq_group = self.swapped[0]
             # If the sequence group has been preempted in this step, stop.
             if seq_group in preempted:

cacheflow/master/server.py

@@ -180,9 +180,9 @@ 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, 32], help='token block size')
+    parser.add_argument('--block-size', type=int, default=16, choices=[1, 2, 4, 8, 16, 32, 64, 128, 256], 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')
+    parser.add_argument('--dtype', type=str, default='half', choices=['half'], 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('--swap-space', type=int, default=20, help='CPU swap space size (GiB) per GPU')

csrc/attention.cpp

@@ -11,25 +11,9 @@ void single_query_cached_kv_attention(
   int block_size,
   int max_context_len);
 
-void multi_query_cached_kv_attention(
-  torch::Tensor& cu_query_lens,
-  torch::Tensor& out,
-  torch::Tensor& query,
-  torch::Tensor& key_cache,
-  torch::Tensor& value_cache,
-  float scale,
-  torch::Tensor& block_tables,
-  torch::Tensor& context_lens,
-  int block_size,
-  int max_context_len);
-
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def(
     "single_query_cached_kv_attention",
     &single_query_cached_kv_attention,
     "Compute the attention between an input query and the cached key/value tensors");
-  m.def(
-    "multi_query_cached_kv_attention",
-    &multi_query_cached_kv_attention,
-    "Compute the attention between multiple input queries and the cached key/value tensors");
 }