wehub-resource-sync
95 lines
2.9 KiB
Python
95 lines
2.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
"""Utility helpers for the TensorRT-LLM integration."""
|
|
|
|
# Standard
|
|
from typing import TYPE_CHECKING
|
|
import os
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
from lmcache.v1.config import LMCacheEngineConfig
|
|
from lmcache.v1.metadata import LMCacheMetadata
|
|
|
|
if TYPE_CHECKING:
|
|
# Third Party
|
|
from tensorrt_llm.llmapi.llm_args import TorchLlmArgs
|
|
|
|
logger = init_logger(__name__)
|
|
|
|
ENGINE_NAME = "trtllm-instance"
|
|
|
|
|
|
def lmcache_get_config() -> LMCacheEngineConfig:
|
|
"""Return an LMCacheEngineConfig from ``LMCACHE_CONFIG_FILE`` or env.
|
|
|
|
Mirrors the in-process pattern used by the vLLM adapter: prefer a
|
|
config file if the env var is set; otherwise pull from individual
|
|
``LMCACHE_*`` environment variables.
|
|
"""
|
|
if "LMCACHE_CONFIG_FILE" not in os.environ:
|
|
logger.warning(
|
|
"No LMCache configuration file is set. Trying to read"
|
|
" configurations from the environment variables."
|
|
)
|
|
logger.warning(
|
|
"You can set the configuration file through "
|
|
"the environment variable: LMCACHE_CONFIG_FILE"
|
|
)
|
|
config = LMCacheEngineConfig.from_env()
|
|
else:
|
|
config_file = os.environ["LMCACHE_CONFIG_FILE"]
|
|
logger.info("Loading LMCache config file %s", config_file)
|
|
config = LMCacheEngineConfig.from_file(config_file)
|
|
|
|
return config
|
|
|
|
|
|
def create_trtllm_metadata(
|
|
llm_args: "TorchLlmArgs",
|
|
kv_cache_tensor: torch.Tensor,
|
|
config: LMCacheEngineConfig,
|
|
num_kv_heads: int,
|
|
head_dim: int,
|
|
) -> LMCacheMetadata:
|
|
"""Construct LMCacheMetadata from TRT-LLM args and the KV pool tensor.
|
|
|
|
Args:
|
|
llm_args: TRT-LLM ``TorchLlmArgs``.
|
|
kv_cache_tensor: TRT-LLM KV pool tensor of shape
|
|
``[num_blocks, num_layers, kv_factor, flat]`` where
|
|
``flat = num_kv_heads * tokens_per_block * head_dim``.
|
|
config: The LMCache engine config (used for ``chunk_size``).
|
|
num_kv_heads: Per-rank number of KV attention heads.
|
|
head_dim: Per-head dimension.
|
|
|
|
Returns:
|
|
Populated :class:`LMCacheMetadata`.
|
|
"""
|
|
# Third Party
|
|
import tensorrt_llm
|
|
|
|
_, num_layers, kv_factor, _ = kv_cache_tensor.shape
|
|
kv_shape = (num_layers, kv_factor, config.chunk_size, num_kv_heads, head_dim)
|
|
|
|
rank = tensorrt_llm.mpi_rank()
|
|
tp_size = llm_args.tensor_parallel_size
|
|
pp_size = llm_args.pipeline_parallel_size
|
|
world_size = tp_size * pp_size
|
|
local_world_size = tp_size
|
|
local_worker_id = rank % local_world_size
|
|
model_name = str(getattr(llm_args, "model", "unknown_model"))
|
|
|
|
return LMCacheMetadata(
|
|
model_name=model_name,
|
|
world_size=world_size,
|
|
local_world_size=local_world_size,
|
|
worker_id=rank,
|
|
local_worker_id=local_worker_id,
|
|
kv_dtype=kv_cache_tensor.dtype,
|
|
kv_shape=kv_shape,
|
|
chunk_size=config.chunk_size,
|
|
)
|