lmcache-cli 0.4.5.dev0__py3-none-any.whl

lmcache/integration/vllm/tests/test_mm_hash_utils.py

@@ -0,0 +1,112 @@
+# SPDX-License-Identifier: Apache-2.0
+# Standard
+import dataclasses
+
+# Third Party
+import torch
+
+# First Party
+from lmcache.integration.vllm.utils import (
+    apply_mm_hashes_to_token_ids,
+    hex_hash_to_int16,
+)
+
+
+@dataclasses.dataclass(frozen=True)
+class DummyPlaceholderRange:
+    offset: int
+    length: int
+
+
+def test_hex_hash_to_int16_accepts_hex_and_non_hex() -> None:
+    # Hex behavior preserved (with and without 0x prefix).
+    assert hex_hash_to_int16("0000") == 0
+    assert hex_hash_to_int16("ffff") == 0xFFFF
+    assert hex_hash_to_int16("0xFFFF") == 0xFFFF
+    assert hex_hash_to_int16("0x0001") == 1
+
+    # Non-hex identifiers must not raise and must be deterministic.
+    s = "chatcmpl-a2a48871c4aad192-image-0"
+    v1 = hex_hash_to_int16(s)
+    v2 = hex_hash_to_int16(s)
+    assert isinstance(v1, int)
+    assert 0 <= v1 <= 0xFFFF
+    assert v1 == v2
+
+
+def test_hex_hash_to_int16_hex_variants_whitespace_and_truncation() -> None:
+    # Whitespace should be ignored and case should not matter.
+    assert hex_hash_to_int16(" FfFf ") == 0xFFFF
+    assert hex_hash_to_int16("\n0x00aB\t") == 0x00AB
+
+    # Long hex should be truncated to 16 bits via masking.
+    assert hex_hash_to_int16("123456") == 0x3456
+    assert hex_hash_to_int16("0x123456") == 0x3456
+
+
+def test_hex_hash_to_int16_empty_and_invalid_hex_are_safe_and_deterministic() -> None:
+    # Empty (or effectively empty) values should not raise.
+    for s in ("", "   ", "0x"):
+        v1 = hex_hash_to_int16(s)
+        v2 = hex_hash_to_int16(s)
+        assert isinstance(v1, int)
+        assert 0 <= v1 <= 0xFFFF
+        assert v1 == v2
+
+    # Invalid "hex-looking" strings must fall back to hashing.
+    for s in ("0xGG", "deadbeeg", "0x12xz"):
+        v1 = hex_hash_to_int16(s)
+        v2 = hex_hash_to_int16(s)
+        assert isinstance(v1, int)
+        assert 0 <= v1 <= 0xFFFF
+        assert v1 == v2
+
+
+def test_hex_hash_to_int16_non_string_inputs_are_safe() -> None:
+    # Be defensive: callers may pass None or other non-string types.
+    for val in (None, 0, 12345, 3.14, b"deadbeef"):
+        v1 = hex_hash_to_int16(val)  # type: ignore[arg-type]
+        v2 = hex_hash_to_int16(val)  # type: ignore[arg-type]
+        assert isinstance(v1, int)
+        assert 0 <= v1 <= 0xFFFF
+        assert v1 == v2
+
+
+def test_apply_mm_hashes_to_token_ids_handles_non_hex_mm_hash() -> None:
+    token_ids = torch.arange(0, 10, dtype=torch.long)
+    mm_hashes = ["chatcmpl-a2a48871c4aad192-image-0"]
+    mm_positions = [DummyPlaceholderRange(offset=2, length=4)]
+
+    out = apply_mm_hashes_to_token_ids(token_ids.clone(), mm_hashes, mm_positions)
+    expected_val = hex_hash_to_int16(mm_hashes[0])
+    assert out[2:6].tolist() == [expected_val] * 4
+
+
+def test_apply_mm_hashes_to_token_ids_out_of_bounds_is_safe() -> None:
+    token_ids = torch.zeros(5, dtype=torch.long)
+    mm_hashes = ["deadbeef"]
+    mm_positions = [DummyPlaceholderRange(offset=999, length=10)]
+
+    out = apply_mm_hashes_to_token_ids(token_ids.clone(), mm_hashes, mm_positions)
+    assert out.tolist() == token_ids.tolist()
+
+
+def test_apply_mm_hashes_to_token_ids_multiple_placeholders_and_length_mismatch() -> (
+    None
+):
+    token_ids = torch.zeros(12, dtype=torch.long)
+    mm_hashes = ["deadbeef", "chatcmpl-a2a48871c4aad192-image-0", "EXTRA_HASH_IGNORED"]
+    mm_positions = [
+        DummyPlaceholderRange(offset=0, length=3),
+        DummyPlaceholderRange(offset=5, length=4),
+    ]
+
+    out = apply_mm_hashes_to_token_ids(token_ids.clone(), mm_hashes, mm_positions)
+    v0 = hex_hash_to_int16(mm_hashes[0])
+    v1 = hex_hash_to_int16(mm_hashes[1])
+
+    assert out[0:3].tolist() == [v0] * 3
+    assert out[5:9].tolist() == [v1] * 4
+    # Other regions remain unchanged.
+    assert out[3:5].tolist() == [0, 0]
+    assert out[9:12].tolist() == [0, 0, 0]

lmcache/integration/vllm/utils.py

@@ -0,0 +1,433 @@
+# 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
+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()
+                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 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 get_vllm_torch_dev():
+    """
+    Returns the torch device and device name for the vLLM engine.
+    e.g. (torch.cuda, "cuda") or (torch.xpu, "xpu")
+    """
+    # Third Party
+    from vllm.platforms import current_platform
+
+    if current_platform.is_cuda_alike():
+        logger.info("CUDA device is available. Using CUDA for LMCache engine.")
+        torch_dev = torch.cuda
+        dev_name = "cuda"
+    elif current_platform.is_xpu():
+        logger.info("XPU device is available. Using XPU for LMCache engine.")
+        torch_dev = torch.xpu
+        dev_name = "xpu"
+    elif hasattr(torch, "hpu") and torch.hpu.is_available():
+        logger.info("HPU device is available. Using HPU for LMCache engine.")
+        torch_dev = torch.hpu
+        dev_name = "hpu"
+    else:
+        raise RuntimeError("Unsupported device platform for LMCache engine.")
+    return torch_dev, dev_name
+
+
+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)
+    """
+    parallel_config = vllm_config.parallel_config
+    global_rank = parallel_config.rank
+    global_world_size = parallel_config.world_size
+    torch_dev, dev_name = get_vllm_torch_dev()
+    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