wehub-resource-sync
119 lines
4.3 KiB
Python
119 lines
4.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from dataclasses import dataclass
|
|
from typing import Optional
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
from lmcache.v1.kv_layer_groups import KVLayerGroupsManager
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class LMCacheMetadata:
|
|
"""
|
|
LMCacheMetadata should be extracted from the northbound
|
|
serving engine configuration and wrap the extraction of
|
|
attributes (e.g. model name, tp rank, etc.)
|
|
"""
|
|
|
|
"""name of the LLM model"""
|
|
model_name: str
|
|
""" global world size when running under a distributed setting
|
|
(total number of workers)"""
|
|
world_size: int
|
|
""" host world size (workers on active localhost)
|
|
This information can be useful for multi-node
|
|
deployment. Will be the same as world_size
|
|
in single-node deployments.
|
|
"""
|
|
local_world_size: int
|
|
""" worker id when running under a distributed setting """
|
|
worker_id: int
|
|
""" host worker id (a gpu bound worker id on active localhost)
|
|
This information can be useful for multi-node deployment.
|
|
Will be the same as worker_id in single-node deployments.
|
|
"""
|
|
local_worker_id: int
|
|
""" the data type of kv tensors """
|
|
# (Deprecated) Will be replaced by kv_layer_groups_manager in the future
|
|
kv_dtype: torch.dtype
|
|
""" the shape of kv tensors """
|
|
# (Deprecated) Will be replaced by kv_layer_groups_manager in the future
|
|
""" (num_layer, 2, chunk_size, num_kv_head, head_size) """
|
|
kv_shape: tuple[int, int, int, int, int]
|
|
""" whether use MLA"""
|
|
use_mla: bool = False
|
|
""" the role of the current instance (e.g., 'scheduler', 'worker') """
|
|
role: Optional[str] = None
|
|
""" the first rank of the distributed setting """
|
|
# TODO(baoloongmao): first_rank should be configurable
|
|
first_rank = 0
|
|
served_model_name: Optional[str] = None
|
|
"""chunk size"""
|
|
chunk_size: int = 256
|
|
"""Manager for groups of layers with identical KV cache structure.
|
|
|
|
``None`` until the serving-engine adapter (e.g. the vLLM connector)
|
|
registers KV caches and constructs the manager via
|
|
:class:`KVLayerGroupsManager`. Consumers must guard against ``None``
|
|
when accessed before registration.
|
|
"""
|
|
kv_layer_groups_manager: Optional[KVLayerGroupsManager] = None
|
|
""" engine_id for RPC path (used by lookup client/server) """
|
|
engine_id: Optional[str] = None
|
|
""" extra config from kv_connector (e.g., lmcache_rpc_port) """
|
|
kv_connector_extra_config: Optional[dict] = None
|
|
|
|
def is_first_rank(self) -> bool:
|
|
"""Check if the current worker is the first rank"""
|
|
return self.worker_id == self.first_rank
|
|
|
|
def get_dtypes(self) -> list[torch.dtype]:
|
|
"""Return per-group dtypes, or the legacy single dtype if the
|
|
manager has not been registered yet (e.g. some unit tests)."""
|
|
klg_manager = self.kv_layer_groups_manager
|
|
if klg_manager is not None and klg_manager.kv_layer_groups:
|
|
return [group.dtype for group in klg_manager.kv_layer_groups]
|
|
return [self.kv_dtype]
|
|
|
|
def get_shapes(self, num_tokens: Optional[int] = None) -> list[torch.Size]:
|
|
"""Get the shapes of the KV cache in LMCache"""
|
|
if num_tokens is None:
|
|
num_tokens = self.chunk_size
|
|
klg_manager = self.kv_layer_groups_manager
|
|
if klg_manager is not None and klg_manager.kv_layer_groups:
|
|
# Read kv_size from each group's shape_desc rather than self.use_mla
|
|
# so heterogeneous groups (should any ever co-exist) are handled.
|
|
return [
|
|
torch.Size(
|
|
[
|
|
group.shape_desc.kv_size,
|
|
group.num_layers,
|
|
num_tokens,
|
|
group.hidden_dim_size,
|
|
]
|
|
)
|
|
for group in klg_manager.kv_layer_groups
|
|
]
|
|
return [
|
|
torch.Size(
|
|
[
|
|
self.kv_shape[1],
|
|
self.kv_shape[0],
|
|
num_tokens,
|
|
self.kv_shape[3] * self.kv_shape[4],
|
|
]
|
|
)
|
|
]
|
|
|
|
def get_num_groups(self) -> int:
|
|
klg_manager = self.kv_layer_groups_manager
|
|
if klg_manager is not None and klg_manager.kv_layer_groups:
|
|
return klg_manager.num_groups
|
|
return 1
|