diff --git a/dlinfer/framework/lmdeploy_ext/device/__init__.py b/dlinfer/framework/lmdeploy_ext/device/__init__.py index 61653bc8..53fc28ea 100644 --- a/dlinfer/framework/lmdeploy_ext/device/__init__.py +++ b/dlinfer/framework/lmdeploy_ext/device/__init__.py @@ -222,47 +222,55 @@ def _cache_engine_allocate_caches( ##### patch state cache engine ##### def patch_state_cache_engine(): - from typing import List, Tuple + from typing import List, Optional, Sequence, Tuple + + from lmdeploy.pytorch.config import StateCacheSpec from lmdeploy.pytorch.engine import cache_engine - from lmdeploy.pytorch.engine.cache_engine import CacheDesc @staticmethod def _state_cache_engine_allocate_caches( num_caches: int, - state_shapes: List[Tuple[Tuple[int], torch.dtype]], + state_shapes: List[Tuple[Tuple[int, ...], torch.dtype]], device: torch.device, + state_specs: Optional[List[StateCacheSpec]] = None, + num_layers: Optional[int] = None, ): """Allocate cache implement. - Each state is allocated as an independent contiguous tensor of shape - (num_caches, *shape). A single shared pool of shape - (num_caches, total_pool_size) would give views with stride[0] == - total_pool_size instead of the per-state numel, making every slice - non-contiguous and breaking NPU ops that require contiguous input. + Each state is allocated as an independent contiguous tensor. A single + shared pool would give state views whose strides include the full pool + row, breaking NPU ops that require contiguous input. Layer-scoped named + caches use (num_rows, num_caches, *shape), matching lmdeploy's logical + layout while keeping every per-layer cache contiguous. """ - cache_dtype = torch.int8 - if len(state_shapes) == 0 or num_caches == 0: + cache_dtype = torch.uint8 + state_specs = state_specs or [] + if (len(state_shapes) == 0 and len(state_specs) == 0) or num_caches == 0: return torch.empty((0, 0), dtype=cache_dtype, device=device), [] - cache_descs = [CacheDesc(shape, dtype) for shape, dtype in state_shapes] + resources = cache_engine.StateCacheEngine._get_state_cache_resources( + state_shapes, state_specs=state_specs, num_layers=num_layers + ) # Allocate each state as a separate contiguous tensor. caches = [] - for desc in cache_descs: - cache = torch.zeros( - (num_caches, *desc.shape), dtype=desc.dtype, device=device - ) + for resource in resources: + desc = resource.desc + cache_shape = (num_caches, *desc.shape) + if resource.layout is not None: + cache_shape = (resource.num_rows, num_caches, *desc.shape[1:]) + cache = torch.zeros(cache_shape, dtype=desc.dtype, device=device) caches.append(cache) # mem_pool is used by two callers: # 1. get_cache_state_size(): always calls with device='meta' to compute byte # counts — the tensor is never materialised on a real device. - # 2. init_caches(): patched below to zero individual caches directly, so it - # no longer touches mem_pool at all. + # 2. init_caches()/copy_caches(): patched below to operate on the independent + # cache tensors directly, so they no longer touch mem_pool at all. # Therefore we only need a correctly-sized pool on 'meta'; for real devices we # return an empty placeholder to avoid doubling the state-cache memory footprint. - total_bytes = sum(desc.aligned_size for desc in cache_descs) + total_bytes = sum(resource.desc.aligned_size for resource in resources) if str(device) == "meta": mem_pool = torch.empty( (num_caches, total_bytes), dtype=cache_dtype, device=device @@ -271,6 +279,14 @@ def _state_cache_engine_allocate_caches( mem_pool = torch.empty(0, dtype=cache_dtype, device=device) return mem_pool, caches + def _state_cache_slot_dim(self, cache_idx: int): + """Return the state-slot dimension for one cache tensor.""" + cache_names = getattr(self, "_state_cache_names", []) + layer_maps = getattr(self, "_state_cache_layer_maps", {}) + if cache_idx < len(cache_names) and cache_names[cache_idx] in layer_maps: + return 1 + return 0 + def _state_cache_engine_init_caches(self, idx: torch.Tensor, mask: torch.Tensor): """Initialize state caches by zeroing each individual cache tensor.""" if idx is None: @@ -280,12 +296,54 @@ def _state_cache_engine_init_caches(self, idx: torch.Tensor, mask: torch.Tensor) num_caches = self.cache_config.num_state_caches cache_masks = torch.zeros((num_caches,), dtype=torch.bool, device=idx.device) cache_masks.index_copy_(0, idx, mask) - for cache in self._state_caches: - reshaped_mask = cache_masks.view((-1,) + (1,) * (cache.dim() - 1)) + for cache_idx, cache in enumerate(self._state_caches): + slot_dim = _state_cache_slot_dim(self, cache_idx) + mask_shape = [1] * cache.dim() + mask_shape[slot_dim] = num_caches + reshaped_mask = cache_masks.view(mask_shape) cache.masked_fill_(reshaped_mask, 0) + def _state_cache_engine_copy_caches( + self, src_idx: int | Sequence[int], dst_idx: int | Sequence[int] + ): + """Copy slots between independently allocated state-cache tensors.""" + if len(self._state_caches) <= 0: + return + + src_list = self._index_list(src_idx) + dst_list = self._index_list(dst_idx) + if len(src_list) != len(dst_list): + raise ValueError( + "src_idx and dst_idx must have the same number of elements." + ) + if len(src_list) == 0: + return + + num_caches = self.cache_config.num_state_caches + self._validate_index_bounds(src_list, num_caches) + self._validate_index_bounds(dst_list, num_caches) + dst_set = set(dst_list) + if len(dst_set) != len(dst_list): + raise ValueError("dst_idx must not contain duplicate entries.") + if not set(src_list).isdisjoint(dst_set): + raise ValueError( + "src_idx and dst_idx must not overlap for stream-ordered state copies." + ) + + for cache_idx, cache in enumerate(self._state_caches): + slot_dim = _state_cache_slot_dim(self, cache_idx) + for src, dst, length in self._copy_ranges(src_list, dst_list): + src_slice = [slice(None)] * cache.dim() + dst_slice = [slice(None)] * cache.dim() + src_slice[slot_dim] = slice(src, src + length) + dst_slice[slot_dim] = slice(dst, dst + length) + cache[tuple(dst_slice)].copy_( + cache[tuple(src_slice)], non_blocking=True + ) + cache_engine.StateCacheEngine.allocate_caches = _state_cache_engine_allocate_caches cache_engine.StateCacheEngine.init_caches = _state_cache_engine_init_caches + cache_engine.StateCacheEngine.copy_caches = _state_cache_engine_copy_caches def patch_gated_delta_net():