wehub-resource-sync
420 lines
16 KiB
Python
420 lines
16 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# Standard
|
|
from typing import TYPE_CHECKING, Literal, Optional, Tuple
|
|
import hashlib
|
|
import os
|
|
import string
|
|
import threading
|
|
|
|
if TYPE_CHECKING:
|
|
from vllm.config import ModelConfig, VllmConfig
|
|
from vllm.multimodal.inputs import PlaceholderRange
|
|
from vllm.v1.request import Request
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.logging import init_logger
|
|
from lmcache.v1.config import LMCacheEngineConfig, load_ec_engine_config
|
|
from lmcache.v1.config_base import apply_remote_configs, fetch_remote_config
|
|
|
|
if TYPE_CHECKING:
|
|
# First Party
|
|
from lmcache.v1.gpu_connector.utils import LayoutHints
|
|
|
|
logger = init_logger(__name__)
|
|
ENGINE_NAME = "vllm-instance"
|
|
|
|
# Thread-safe singleton storage
|
|
_config_instance: Optional[LMCacheEngineConfig] = None
|
|
_config_lock = threading.Lock()
|
|
|
|
|
|
def is_false(value: str) -> bool:
|
|
"""Check if the given string value is equivalent to 'false'."""
|
|
return value.lower() in ("false", "0", "no", "n", "off")
|
|
|
|
|
|
def vllm_layout_hints() -> "LayoutHints":
|
|
"""Build layout_hints dict by querying vLLM at runtime."""
|
|
hints: dict[str, str] = {}
|
|
kv_layout = try_get_vllm_kv_cache_layout()
|
|
if kv_layout is not None:
|
|
hints["kv_layout"] = kv_layout
|
|
return hints # type: ignore[return-value]
|
|
|
|
|
|
def try_get_vllm_kv_cache_layout() -> Literal["NHD", "HND"] | None:
|
|
"""Try to query the KV cache layout from vLLM at runtime.
|
|
|
|
Returns ``"NHD"`` or ``"HND"`` if vLLM is available and the layout
|
|
has been configured, otherwise ``None``.
|
|
|
|
Please only call this where vllm is available (i.e. not in the MP server)
|
|
We will print an error if we try to get vllm kv layout where vllm
|
|
is not available.
|
|
"""
|
|
|
|
# Third Party
|
|
try:
|
|
# Third Party
|
|
from vllm.v1.attention.backends.utils import ( # type: ignore[import-untyped]
|
|
get_kv_cache_layout,
|
|
)
|
|
|
|
return get_kv_cache_layout()
|
|
except Exception:
|
|
logger.error(
|
|
"vLLM is not available but tried to query kv cache "
|
|
"layout information, cannot get KV cache layout"
|
|
)
|
|
return None
|
|
|
|
|
|
def lmcache_get_or_create_config() -> LMCacheEngineConfig:
|
|
"""Get the LMCache configuration from the environment variable
|
|
`LMCACHE_CONFIG_FILE`. If the environment variable is not set, this
|
|
function will return the default configuration.
|
|
|
|
This function is thread-safe and implements singleton pattern,
|
|
ensuring the configuration is loaded only once.
|
|
|
|
After loading the configuration, if 'remote_config_url' is configured,
|
|
this function will attempt to fetch additional configuration from the
|
|
remote config service. The current config and LMCACHE environment
|
|
variables will be sent to the service, along with 'appid' if set.
|
|
"""
|
|
global _config_instance
|
|
|
|
# Double-checked locking for thread-safe singleton
|
|
if _config_instance is None:
|
|
with _config_lock:
|
|
if _config_instance is None: # Check again within lock
|
|
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_instance = LMCacheEngineConfig.from_env()
|
|
# from_env() doesn't call validate(); the file path
|
|
# gets it via update_config_from_env() below, but the
|
|
# env-only path needs an explicit call.
|
|
_config_instance.validate()
|
|
else:
|
|
config_file = os.environ["LMCACHE_CONFIG_FILE"]
|
|
logger.info(f"Loading LMCache config file {config_file}")
|
|
_config_instance = LMCacheEngineConfig.from_file(config_file)
|
|
# Update config from environment variables
|
|
_config_instance.update_config_from_env()
|
|
|
|
# Fetch and apply remote configuration if configured
|
|
remote_config_url = _config_instance.remote_config_url
|
|
if remote_config_url:
|
|
logger.info(
|
|
"Fetching remote configuration from %s", remote_config_url
|
|
)
|
|
app_id = _config_instance.app_id
|
|
remote_response = fetch_remote_config(
|
|
remote_config_url, app_id, _config_instance
|
|
)
|
|
if remote_response:
|
|
_config_instance = apply_remote_configs(
|
|
_config_instance, remote_response
|
|
)
|
|
else:
|
|
logger.warning(
|
|
"Failed to fetch remote configuration from %s. "
|
|
"Using local configuration only.",
|
|
remote_config_url,
|
|
)
|
|
return _config_instance
|
|
|
|
|
|
def create_lmcache_ec_config() -> LMCacheEngineConfig:
|
|
"""Create EC config from LMCache config plus EC-specific overrides."""
|
|
return load_ec_engine_config(base_config=lmcache_get_or_create_config())
|
|
|
|
|
|
def hex_hash_to_int16(s: str) -> int:
|
|
"""
|
|
Convert a hash identifier into a 16-bit integer.
|
|
|
|
Historically, LMCache expected multimodal identifiers to be hex strings.
|
|
In practice (e.g., OpenAI-style multimodal requests), identifiers may be
|
|
arbitrary strings like `chatcmpl-...-image-0`. This function therefore:
|
|
- Parses hex strings (optionally prefixed with `0x`) as before, or
|
|
- Falls back to a stable string hash (SHA-256) when the input is not hex.
|
|
"""
|
|
# Be defensive: vLLM may pass non-string identifiers.
|
|
s = "" if s is None else str(s)
|
|
s_stripped = s.strip()
|
|
|
|
# Fast-path: pure hex (optionally 0x-prefixed).
|
|
hex_part = s_stripped[2:] if s_stripped.lower().startswith("0x") else s_stripped
|
|
if hex_part and all(c in string.hexdigits for c in hex_part):
|
|
try:
|
|
return int(hex_part, 16) & 0xFFFF
|
|
except ValueError:
|
|
# Extremely unlikely (e.g., oversized/odd formatting); fall back to hashing.
|
|
pass
|
|
|
|
# Fallback: stable 16-bit value derived from the full identifier string.
|
|
digest = hashlib.sha256(s_stripped.encode("utf-8")).digest()
|
|
return int.from_bytes(digest[:2], byteorder="big", signed=False)
|
|
|
|
|
|
def apply_mm_hashes_to_token_ids(
|
|
token_ids: torch.Tensor,
|
|
mm_hashes: list[str],
|
|
mm_positions: list["PlaceholderRange"],
|
|
) -> torch.Tensor:
|
|
"""
|
|
Overwrite token_ids in-place for multimodal placeholders using
|
|
efficient slice assignments.
|
|
"""
|
|
n = token_ids.size(0)
|
|
for hash_str, placeholder in zip(mm_hashes, mm_positions, strict=False):
|
|
start, length = placeholder.offset, placeholder.length
|
|
if start >= n:
|
|
continue
|
|
end = min(start + length, n)
|
|
token_ids[start:end] = hex_hash_to_int16(hash_str)
|
|
return token_ids
|
|
|
|
|
|
def mla_enabled(model_config: "ModelConfig") -> bool:
|
|
return (
|
|
hasattr(model_config, "use_mla")
|
|
and isinstance(model_config.use_mla, bool)
|
|
and model_config.use_mla
|
|
)
|
|
|
|
|
|
def create_lmcache_metadata(
|
|
vllm_config=None,
|
|
model_config=None,
|
|
parallel_config=None,
|
|
cache_config=None,
|
|
role=None,
|
|
):
|
|
"""
|
|
Create LMCacheMetadata from vLLM configuration.
|
|
|
|
This function extracts common metadata creation logic that was duplicated
|
|
across multiple files.
|
|
|
|
Args:
|
|
vllm_config: vLLM configuration object containing model, parallel, and
|
|
cache configs (alternative to individual config parameters)
|
|
model_config: Model configuration (alternative to vllm_config)
|
|
parallel_config: Parallel configuration (alternative to vllm_config)
|
|
cache_config: Cache configuration (alternative to vllm_config)
|
|
|
|
Returns:
|
|
tuple: (LMCacheMetadata, LMCacheEngineConfig)
|
|
"""
|
|
# Third Party
|
|
# Try to import from old location before merged https://github.com/vllm-project/vllm/pull/26908
|
|
try:
|
|
# Third Party
|
|
from vllm.utils.torch_utils import get_kv_cache_torch_dtype
|
|
except ImportError:
|
|
# Third Party
|
|
from vllm.utils import get_kv_cache_torch_dtype
|
|
# First Party
|
|
from lmcache.v1.metadata import LMCacheMetadata
|
|
|
|
config = lmcache_get_or_create_config()
|
|
# Support both vllm_config object and individual config parameters
|
|
if vllm_config is not None:
|
|
model_cfg = vllm_config.model_config
|
|
parallel_cfg = vllm_config.parallel_config
|
|
cache_cfg = vllm_config.cache_config
|
|
else:
|
|
model_cfg = model_config
|
|
parallel_cfg = parallel_config
|
|
cache_cfg = cache_config
|
|
|
|
# Get KV cache dtype
|
|
kv_dtype = get_kv_cache_torch_dtype(cache_cfg.cache_dtype, model_cfg.dtype)
|
|
|
|
# Check if MLA is enabled
|
|
use_mla = mla_enabled(model_cfg)
|
|
|
|
# Construct KV shape (for memory pool)
|
|
num_layer = model_cfg.get_num_layers(parallel_cfg)
|
|
chunk_size = config.chunk_size
|
|
num_kv_head = model_cfg.get_num_kv_heads(parallel_cfg)
|
|
head_size = model_cfg.get_head_size()
|
|
kv_shape = (num_layer, 1 if use_mla else 2, chunk_size, num_kv_head, head_size)
|
|
|
|
# Extract engine_id and kv_connector_extra_config from vllm_config if available
|
|
engine_id = None
|
|
kv_connector_extra_config = None
|
|
if vllm_config is not None and hasattr(vllm_config, "kv_transfer_config"):
|
|
kv_transfer_config = vllm_config.kv_transfer_config
|
|
if kv_transfer_config is not None:
|
|
engine_id = getattr(kv_transfer_config, "engine_id", None)
|
|
kv_connector_extra_config = getattr(
|
|
kv_transfer_config, "kv_connector_extra_config", None
|
|
)
|
|
|
|
# Create metadata
|
|
metadata = LMCacheMetadata(
|
|
model_name=model_cfg.model,
|
|
world_size=parallel_cfg.world_size,
|
|
local_world_size=parallel_cfg.world_size,
|
|
worker_id=parallel_cfg.rank,
|
|
local_worker_id=parallel_cfg.rank,
|
|
kv_dtype=kv_dtype,
|
|
kv_shape=kv_shape,
|
|
use_mla=use_mla,
|
|
role=role,
|
|
served_model_name=model_cfg.served_model_name,
|
|
engine_id=engine_id,
|
|
kv_connector_extra_config=kv_connector_extra_config,
|
|
)
|
|
|
|
return metadata, config
|
|
|
|
|
|
def extract_mm_features(
|
|
request: "Request", modify: bool = False
|
|
) -> Tuple[list[str], list["PlaceholderRange"]]:
|
|
"""
|
|
Normalize multimodal information from a Request into parallel lists.
|
|
|
|
This helper reads either:
|
|
1) `request.mm_features` (objects each exposing `.identifier` and
|
|
`.mm_position`), or
|
|
2) legacy fields `request.mm_hashes` and `request.mm_positions`.
|
|
|
|
It returns two equally sized lists: the multimodal hash identifiers and their
|
|
corresponding positions. If the request contains no multimodal info, it returns
|
|
`([], [])`.
|
|
|
|
Args:
|
|
request (Request): The source object.
|
|
modify (bool):
|
|
Controls copy semantics for the legacy-path return values.
|
|
- If True and legacy fields are used, shallow-copies are returned so
|
|
the caller can mutate the lists without affecting `request`.
|
|
- If False, the original legacy sequences are returned as-is
|
|
(zero-copy); treat them as read-only.
|
|
|
|
Returns:
|
|
Tuple[list[str], list[PlaceholderRange]]: (`mm_hashes`, `mm_positions`).
|
|
May be `([], [])` when no multimodal data is present.
|
|
"""
|
|
if getattr(request, "mm_features", None):
|
|
mm_hashes, mm_positions = zip(
|
|
*((f.identifier, f.mm_position) for f in request.mm_features), strict=False
|
|
)
|
|
return (list(mm_hashes), list(mm_positions))
|
|
elif getattr(request, "mm_hashes", None):
|
|
if modify:
|
|
return (request.mm_hashes.copy(), request.mm_positions.copy())
|
|
else:
|
|
return (request.mm_hashes, request.mm_positions)
|
|
else:
|
|
return ([], [])
|
|
|
|
|
|
def get_size_bytes(shapes: list[torch.Size], kv_dtypes: list[torch.dtype]):
|
|
"""
|
|
Calculate the size in bytes with the given shapes and dtypes.
|
|
"""
|
|
assert len(shapes) == len(kv_dtypes), (
|
|
f"shapes and dtypes must have the same length, "
|
|
f"but got {len(shapes)} and {len(kv_dtypes)}"
|
|
)
|
|
return sum(
|
|
shape.numel() * kv_dtype.itemsize
|
|
for shape, kv_dtype in zip(shapes, kv_dtypes, strict=True)
|
|
)
|
|
|
|
|
|
def calculate_local_rank_and_world_size(vllm_config: "VllmConfig") -> Tuple[int, int]:
|
|
"""
|
|
Calculate the local worker id and local world size.
|
|
|
|
Current assumption (TODO: add custom logic in the future):
|
|
- Tensor Parallel is intra-node
|
|
- Pipeline Parallel is inter-node
|
|
|
|
Returns:
|
|
Tuple[int, int]: (local_worker_id, local_world_size)
|
|
"""
|
|
# First Party
|
|
from lmcache import torch_dev
|
|
|
|
parallel_config = vllm_config.parallel_config
|
|
global_rank = parallel_config.rank
|
|
global_world_size = parallel_config.world_size
|
|
num_gpus = torch_dev.device_count()
|
|
if global_world_size <= num_gpus:
|
|
# single node case
|
|
return parallel_config.rank, parallel_config.world_size
|
|
else:
|
|
tp_size = parallel_config.tensor_parallel_size
|
|
pp_size = parallel_config.pipeline_parallel_size
|
|
local_world_size = global_world_size // pp_size
|
|
assert local_world_size == tp_size, (
|
|
"LMCache is operating under the assumption that the "
|
|
"local world size is equal to the tensor parallel size "
|
|
"in multi-node deployment."
|
|
)
|
|
local_worker_id = global_rank % local_world_size
|
|
return local_worker_id, local_world_size
|
|
|
|
|
|
def validate_mla_config(config: LMCacheEngineConfig, use_mla: bool) -> None:
|
|
"""Validate MLA-related configuration."""
|
|
if use_mla and (config.remote_serde != "naive" and config.remote_serde is not None):
|
|
raise ValueError("MLA only works with naive serde mode..")
|
|
|
|
if use_mla and config.use_layerwise and config.enable_blending:
|
|
raise ValueError(
|
|
"We haven't supported MLA with Cacheblend yet. Please disable blending."
|
|
)
|
|
|
|
|
|
def calculate_draft_layers(vllm_config: "VllmConfig") -> int:
|
|
"""Calculate the number of draft layers for speculative decoding."""
|
|
assert vllm_config is not None, "vllm_config required for vLLM mode"
|
|
|
|
num_draft_layers = 0
|
|
model_config = vllm_config.model_config
|
|
|
|
if vllm_config.speculative_config is not None:
|
|
logger.info(
|
|
"vllm_config.speculative_config: %s", vllm_config.speculative_config
|
|
)
|
|
if vllm_config.speculative_config.method == "deepseek_mtp":
|
|
num_draft_layers = getattr(
|
|
model_config.hf_config, "num_nextn_predict_layers", 0
|
|
)
|
|
elif vllm_config.speculative_config.use_eagle():
|
|
try:
|
|
draft_model_config = vllm_config.speculative_config.draft_model_config
|
|
num_draft_layers = draft_model_config.get_num_layers(
|
|
vllm_config.parallel_config
|
|
)
|
|
logger.info("EAGLE detected %d extra layer(s)", num_draft_layers)
|
|
except Exception:
|
|
logger.info(
|
|
"EAGLE detected, but failed to get the number of extra layers"
|
|
"falling back to 1"
|
|
)
|
|
num_draft_layers = 1
|
|
return num_draft_layers
|
|
|
|
|
|
def is_dp_rank0(vllm_config: "VllmConfig") -> bool:
|
|
return vllm_config.parallel_config.data_parallel_rank_local == 0
|