lmdeploy support kernel block size (#4421)

* add kernel block config

* support kernel block size

* add comment

* fix

* fix format

* move cal kernel offs to _tensorlize_block_offsets

* fix

* add map kernel offs func

---------

Co-authored-by: maoruihan <maoruihan@stonewise.cn>

Author: Tsundoku958

lmdeploy/cli/serve.py

@@ -106,6 +106,7 @@ def add_parser_api_server():
         ArgumentHelper.dllm_confidence_threshold(pt_group)
         ArgumentHelper.enable_return_routed_experts(pt_group)
         ArgumentHelper.distributed_executor_backend(pt_group)
+        ArgumentHelper.kernel_block_size(pt_group)
 
         # common engine args
         dtype_act = ArgumentHelper.dtype(pt_group)
@@ -226,6 +227,7 @@ def api_server(args):
                 max_batch_size=max_batch_size,
                 cache_max_entry_count=args.cache_max_entry_count,
                 block_size=args.cache_block_seq_len,
+                kernel_block_size=args.kernel_block_size,
                 session_len=args.session_len,
                 adapters=adapters,
                 enable_prefix_caching=args.enable_prefix_caching,

lmdeploy/cli/utils.py

@@ -546,6 +546,19 @@ def cache_block_seq_len(parser):
                                    'if Lora Adapter is specified, this parameter will '
                                    'be ignored')
 
+    @staticmethod
+    def kernel_block_size(parser):
+        """Add argument kernel_block_size to parser."""
+
+        return parser.add_argument('--kernel-block-size',
+                                   type=int,
+                                   default=-1,
+                                   help='The length of the token sequence in a k/v block for kernels. '
+                                   'Only supported by Pytorch Engine. '
+                                   'When set to a different value than --cache-block-seq-len, '
+                                   'memory allocators and prefix cache use --cache-block-seq-len '
+                                   'as the block size, while kernels use --kernel-block-size.')
+
     @staticmethod
     def enable_prefix_caching(parser):
         """Add argument enable_prefix_caching to parser."""

lmdeploy/messages.py

@@ -392,6 +392,7 @@ class PytorchEngineConfig:
     cache_max_entry_count: float = 0.8
     prefill_interval: int = 16
     block_size: int = 64
+    kernel_block_size: int = -1
     num_cpu_blocks: int = 0
     num_gpu_blocks: int = 0
     adapters: dict[str, str] = None
@@ -430,6 +431,8 @@ class PytorchEngineConfig:
 
     def __post_init__(self):
         """Check input validation."""
+        if self.kernel_block_size == -1:
+            self.kernel_block_size = self.block_size
         assert self.dtype in ['auto', 'float16', 'bfloat16']
         assert self.tp >= 1, 'invalid tp'
         assert self.dp >= 1, 'invalid dp'
@@ -442,8 +445,14 @@ def __post_init__(self):
         assert self.num_gpu_blocks >= 0, 'invalid num_gpu_blocks'
         assert self.quant_policy in (0, 4, 8), 'invalid quant_policy'
         assert self.device_type in ['cuda', 'ascend', 'maca', 'camb'], (f'invalid device_type: {self.device_type}')
-        assert self.block_size >= 16 and (self.block_size & (self.block_size - 1)) == 0, \
-            f'block_size must be >= 16 and a power of 2, but got {self.block_size}'
+        assert self.kernel_block_size >= 16 and \
+               (self.kernel_block_size & (self.kernel_block_size - 1)) == 0, \
+               f'kernel_block_size must be >= 16 and a power of 2, but got {self.kernel_block_size}'
+        assert self.block_size >= self.kernel_block_size and \
+               self.block_size % self.kernel_block_size == 0, \
+               (f'block_size must be >= kernel_block_size and an integer multiple '
+                f'of kernel_block_size, but got block_size {self.block_size} '
+                f'and kernel_block_size {self.kernel_block_size}')
         if self.quant_policy > 0 and self.device_type not in ['cuda', 'ascend']:
             assert False, \
                    'kv cache quantization only works for CUDA and ASCEND.'

lmdeploy/pytorch/config.py

@@ -94,6 +94,7 @@ class CacheConfig:
     block_size: int
     num_cpu_blocks: int
     num_gpu_blocks: int
+    kernel_block_size: int = -1
     window_size: int = -1
     cache_max_entry_count: float = 0.8
     max_prefill_token_num: int = 4096
@@ -115,6 +116,8 @@ def __post_init__(self):
         if self.window_size > 1 and self.enable_prefix_caching:
             logger.warning('Prefix caching is not available for window attention.')
             self.enable_prefix_caching = False
+        if self.kernel_block_size == -1:
+            self.kernel_block_size = self.block_size
 
 
 class TPMode(enum.Enum):

lmdeploy/pytorch/engine/cache_engine.py

@@ -84,7 +84,7 @@ def __init__(
         self.cache_config = cache_config
         self.model_config = model_config
 
-        self.block_size = cache_config.block_size
+        self.block_size = cache_config.kernel_block_size
         self.num_layers = model_config.num_layers
         self.kv_cache_dtype = _get_kv_cache_dtype(self.model_config)
 
@@ -198,7 +198,7 @@ def get_k_cache_desc(cls, model_config: ModelConfig, cache_config: CacheConfig,
             head_size = model_config.head_dim
         shape = cls._get_key_block_shape_impl(
             model_config,
-            block_size=cache_config.block_size,
+            block_size=cache_config.kernel_block_size,
             head_size=head_size,
             world_size=world_size,
             quant_policy=cache_config.quant_policy,
@@ -217,7 +217,7 @@ def get_v_cache_desc(cls, model_config: ModelConfig, cache_config: CacheConfig,
             head_size = model_config.head_dim
         shape = cls._get_value_block_shape_impl(
             model_config,
-            block_size=cache_config.block_size,
+            block_size=cache_config.kernel_block_size,
             head_size=head_size,
             world_size=world_size,
             quant_policy=cache_config.quant_policy,
@@ -248,7 +248,7 @@ def get_custom_cache_descs(cls, model_config: ModelConfig, cache_config: CacheCo
         if len(model_config.cache_shapes) == 0:
             return []
 
-        block_size = cache_config.block_size
+        block_size = cache_config.kernel_block_size
 
         descs = []
         for shape, dtype in model_config.cache_shapes:
@@ -263,6 +263,8 @@ def allocate_caches(cls, num_blocks: int, model_config: ModelConfig, cache_confi
         """Allocate caches."""
 
         num_layers = model_config.num_layers
+        kernel_blocks_per_kv = cache_config.block_size // cache_config.kernel_block_size
+        num_blocks *= kernel_blocks_per_kv
 
         # get all descs
         k_cache_desc = cls.get_k_cache_desc(model_config, cache_config, world_size)

lmdeploy/pytorch/engine/config_builder.py

@@ -62,6 +62,7 @@ def build_cache_config(engine_config: PytorchEngineConfig):
         cache_config = CacheConfig(
             max_batches=engine_config.max_batch_size,
             block_size=engine_config.block_size,
+            kernel_block_size=engine_config.kernel_block_size,
             num_cpu_blocks=engine_config.num_cpu_blocks,
             num_gpu_blocks=engine_config.num_gpu_blocks,
             cache_max_entry_count=engine_config.cache_max_entry_count,

lmdeploy/pytorch/engine/inputs_maker.py

@@ -214,6 +214,9 @@ def __init__(
         self.adapter_manager = adapter_manager
         self.config = config
         self.spec_decoding = config.spec_decoding
+        self.cache_config = scheduler.cache_config
+        self.kernel_blocks_per_kv = self.cache_config.block_size // self.cache_config.kernel_block_size
+        self.kernel_block_arange = torch.arange(self.kernel_blocks_per_kv, dtype=self.torch_int_dtype)
 
         # strategies
         self.engine_strategy = engine_strategy
@@ -322,6 +325,29 @@ def _set_adapter_ids(self, model_inputs: ModelInputs, messages: 'SeqList'):
         local_adapter_ids = model_inputs.seq_length.new_tensor(local_adapter_ids)
         model_inputs.local_adapter_ids = local_adapter_ids
 
+    def _map_to_kernel_block_offsets(self, block_offsets: torch.Tensor):
+        """Converts manager block_offsets to kernel block_offsets.
+
+        Example:
+
+            # block_manager block size: 32 tokens,
+            # Kernel block size: 16 tokens
+            # kernel_blocks_per_kv = 2
+            >>> block_manager block offsets = [0, 1, 3]
+            >>> Result kernel block offsets = [0, 1, 2, 3, 6, 7]
+
+            # Each block_manager block id maps to 2 kernel block id:
+            # block_manager block id 0 -> kernel block id [0, 1]
+            # block_manager block id 1 -> kernel block id [2, 3]
+            # block_manager block id 3 -> kernel block id [6, 7]
+        """
+        if self.kernel_blocks_per_kv == 1:
+            return block_offsets
+        batch_size = block_offsets.shape[0]
+        block_offsets = (block_offsets[:, :, None] * self.kernel_blocks_per_kv +
+                         self.kernel_block_arange[None, None, :]).reshape(batch_size, -1)
+        return block_offsets
+
     @torch.inference_mode()
     @record_function('create_model_inputs')
     def create_model_inputs(self, messages: 'SeqList', is_prefill: bool):
@@ -355,6 +381,7 @@ def create_model_inputs(self, messages: 'SeqList', is_prefill: bool):
         # block offsets
         block_offsets = self.scheduler.get_block_tables(messages)
         block_offsets = _tensorlize_block_offsets(block_offsets, dtype=self.torch_int_dtype)
+        block_offsets = self._map_to_kernel_block_offsets(block_offsets)
 
         # num_ignored_history
         num_ignored_history = torch.tensor([msg.num_ignored_history for msg in messages])
@@ -410,6 +437,7 @@ def create_model_inputs_long_context(self,
         # block offsets
         block_offsets = self.scheduler.get_block_tables([seq])
         block_offsets = torch.as_tensor(block_offsets[0], dtype=self.torch_int_dtype)[None]
+        block_offsets = self._map_to_kernel_block_offsets(block_offsets)
 
         # num_ignored_history
         num_ignored_history = torch.tensor([seq.num_ignored_history])
@@ -482,6 +510,7 @@ def create_model_inputs_delta(self):
         # block offsets
         block_offsets = self.scheduler.get_block_tables(valid_seqs)
         block_offsets = _tensorlize_block_offsets(block_offsets, dtype=self.torch_int_dtype)
+        block_offsets = self._map_to_kernel_block_offsets(block_offsets)
 
         # sliding window
         if self.scheduler.cache_config.window_size > 0:

lmdeploy/pytorch/paging/scheduler.py

@@ -353,7 +353,7 @@ def has_unfinished(self):
         return self.has_ready() or self.has_waiting() or self.has_migration_done()
 
     def get_block_tables(self, seqs: SeqList):
-        """Get block table of the sequences."""
+        """Get block tables for the sequences."""
         return [self.block_manager.get_block_table(seq) for seq in seqs]
 
     def evict_seqs(self, running: SeqList):