lmcache-cli 0.5.7rc2.dev0__py3-none-any.whl → 0.5.7rc3.dev0__py3-none-any.whl

lmcache/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.5.7rc2.dev0'
-__version_tuple__ = version_tuple = (0, 5, 7, 'rc2', 'dev0')
+__version__ = version = '0.5.7rc3.dev0'
+__version_tuple__ = version_tuple = (0, 5, 7, 'rc3', 'dev0')
 
-__commit_id__ = commit_id = 'gaa12d34e3'
+__commit_id__ = commit_id = 'gbf85b41fa'

lmcache/cli/commands/bench/server_bench/__init__.py

@@ -22,22 +22,38 @@ class ServerBenchCommand(BaseCommand):
         return "End-to-end test for LMCache MP cache server (GPU mode)."
 
     def add_arguments(self, parser: argparse.ArgumentParser) -> None:
+        """Register benchmark arguments even without the full runtime.
+
+        Args:
+            parser: The argument parser for the server benchmark subcommand.
+        """
         # First Party
-        from lmcache.cli.commands.bench.server_bench.command import (
+        from lmcache.cli.commands.bench.server_bench.arguments import (
             add_server_arguments,
         )
 
         add_server_arguments(parser)
 
     def execute(self, args: argparse.Namespace) -> None:
+        """Load and execute the benchmark, reporting runtime import failures."""
         # Standard
         import copy
         import sys
 
-        # First Party
-        from lmcache.cli.commands.bench.server_bench.command import (
-            run_server_bench,
-        )
+        try:
+            # First Party
+            from lmcache.cli.commands.bench.server_bench.command import (
+                run_server_bench,
+            )
+        except ImportError as exc:
+            print(
+                "ERROR: Unable to load `lmcache bench server`. "
+                "This command requires the full LMCache package.\n"
+                "  Install with: pip install lmcache\n"
+                f"  Original import error: {exc}",
+                file=sys.stderr,
+            )
+            sys.exit(2)
 
         # A comma-separated --flamegraph-mode drives the load once per mode
         # (the modes cannot share a recording window), each rendered to its

lmcache/cli/commands/bench/server_bench/allocator.py

@@ -0,0 +1,78 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Request-exclusive block allocation shared by all TP workers."""
+
+# Standard
+from dataclasses import dataclass
+from typing import Literal
+import random
+
+
+@dataclass(frozen=True)
+class GroupBlocks:
+    """Physical block IDs in logical order for one engine group."""
+
+    group_id: int
+    block_ids: tuple[int, ...]
+
+
+@dataclass(frozen=True)
+class RequestBlocks:
+    """Immutable group allocations shared by every TP worker."""
+
+    groups: tuple[GroupBlocks, ...]
+
+    def as_block_ids(self) -> list[list[int]]:
+        """Return group-ordered IDs for buffer and transfer APIs."""
+        return [list(group.block_ids) for group in self.groups]
+
+
+class BlockAllocator:
+    """Single-threaded, request-exclusive allocation without prefix sharing."""
+
+    def __init__(
+        self,
+        groups: int,
+        capacity: int,
+        *,
+        allocation_order: Literal["sequential", "random"] = "sequential",
+        seed: int = 0,
+    ) -> None:
+        """Create one free pool per group with a reproducible allocation order."""
+        if groups < 0 or capacity < 0:
+            raise ValueError("groups and capacity must be non-negative")
+        if allocation_order not in ("sequential", "random"):
+            raise ValueError(f"unsupported allocation order: {allocation_order}")
+        self._free = [set(range(capacity)) for _ in range(groups)]
+        self._allocation_order = allocation_order
+        self._random = random.Random(seed)
+
+    def allocate(self, count: int) -> RequestBlocks:
+        """Atomically reserve count > 0 exclusive blocks in every group."""
+        if count <= 0:
+            raise ValueError("block count must be positive")
+        for group_id, free in enumerate(self._free):
+            if len(free) < count:
+                raise RuntimeError(
+                    f"group {group_id} needs {count} blocks, only {len(free)} free"
+                )
+        allocated = RequestBlocks(
+            tuple(
+                GroupBlocks(
+                    group_id,
+                    tuple(
+                        self._random.sample(sorted(free), count)
+                        if self._allocation_order == "random"
+                        else sorted(free)[:count]
+                    ),
+                )
+                for group_id, free in enumerate(self._free)
+            )
+        )
+        for group in allocated.groups:
+            self._free[group.group_id].difference_update(group.block_ids)
+        return allocated
+
+    def release(self, allocated: RequestBlocks) -> None:
+        """Return owned blocks after all transfers using them have finished."""
+        for group in allocated.groups:
+            self._free[group.group_id].update(group.block_ids)

lmcache/cli/commands/bench/server_bench/arguments.py

@@ -0,0 +1,226 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Runtime-independent argument registration for ``lmcache bench server``.
+
+Keep this module free of torch, zmq, and benchmark runtime imports so the
+lightweight CLI can show help and parse options before reporting missing
+runtime dependencies. The dtype help is checked against DTYPE_MAP in tests.
+"""
+
+# Standard
+import argparse
+
+_DEFAULT_SHAPE_SPEC = "(2,1024,16,8,128):float16:32"
+
+
+# ---------------------------------------------------------------------------
+# Parser registration
+# ---------------------------------------------------------------------------
+
+
+def add_server_arguments(parser: argparse.ArgumentParser) -> None:
+    """Add ``lmcache bench server`` arguments without loading runtime dependencies.
+
+    Args:
+        parser: The ``ArgumentParser`` for the server bench subcommand.
+    """
+
+    parser.add_argument(
+        "--rpc-url",
+        default="tcp://localhost:5555",
+        help=("MP request endpoint (default: tcp://localhost:5555)"),
+    )
+    parser.add_argument(
+        "--mode",
+        choices=["cpu", "gpu"],
+        default="gpu",
+        help=(
+            "Run mode (default: gpu). In cpu mode the client allocates "
+            "regular CPU KV tensors; the selected TransferContext owns "
+            "any SHM migration or staging."
+        ),
+    )
+    parser.add_argument(
+        "--transfer-mode",
+        choices=["auto", "engine_driven", "lmcache_driven"],
+        default="auto",
+        help=(
+            "Transport routing for STORE/RETRIEVE (default: auto). "
+            "`lmcache_driven` forces the server-driven handle path "
+            "(REGISTER_KV_CACHE + STORE/RETRIEVE), which supports "
+            "both CUDA IPC and CPU SHM for zero-copy transfers. "
+            "`engine_driven` forces the worker-side gather/scatter "
+            "data path (REGISTER_KV_CACHE_ENGINE_DRIVEN_CONTEXT + "
+            "PREPARE/COMMIT). "
+            "`auto` keeps the historical mapping: "
+            "gpu->lmcache_driven, cpu->engine_driven."
+        ),
+    )
+    parser.add_argument(
+        "--tp-size",
+        type=int,
+        default=1,
+        help=(
+            "Simulated tensor-parallel world size (default: 1). Each "
+            "rank registers its own KV cache under a distinct "
+            "instance_id, and STORE / RETRIEVE fan out per rank the "
+            "same way LMCacheMPWorkerAdapter routes them in a real "
+            "vLLM deployment (MLA -> only rank 0 stores; non-MLA -> "
+            "every rank stores; every rank always retrieves)."
+        ),
+    )
+    parser.add_argument(
+        "--use-mla",
+        action="store_true",
+        default=False,
+        help=(
+            "MLA mode: fold all TP ranks into a single kv_worker "
+            "so only rank 0 writes KV and every rank retrieves "
+            "the shared KV object (default: False). "
+            "Also implied when --kvcache-shape-spec declares "
+            "kv_size=1."
+        ),
+    )
+    parser.add_argument(
+        "--num-tokens",
+        type=int,
+        default=512,
+        help="Tokens per request (default: 512)",
+    )
+
+    parser.add_argument(
+        "--block-allocation",
+        choices=["sequential", "random"],
+        default="sequential",
+        help="Block allocation order (default: sequential).",
+    )
+
+    # -- KV cache shape --
+    kv = parser.add_argument_group("KV cache shape")
+    kv.add_argument(
+        "--kvcache-shape-spec",
+        type=str,
+        default=_DEFAULT_SHAPE_SPEC,
+        help=(
+            "KV shape spec. Describes one or more KV layer groups "
+            "separated by ';'. "
+            "Grammar: "
+            "'(kv_size,NB,BS,NH,HS):dtype:layers[;(...):dtype:layers...]'. "
+            "Fields: kv_size=2 for classical K/V or 1 for MLA, "
+            "NB=num_blocks, BS=block_size (tokens/block), "
+            "NH=num_heads, HS=head_size (elements). "
+            "dtype is the element dtype (supported: float16, float32, bfloat16, "
+            "uint8, int8); 'uint8' "
+            "is used for FP8-quantized KV. 'layers' is the number "
+            "of consecutive layers sharing this group's geometry. "
+            "Multi-group example (MLA + classical attention): "
+            "'(1,1024,16,1,128):float16:4;"
+            "(2,1024,16,8,128):float16:28'. "
+            "All groups must share the same NB and BS. "
+            "See lmcache.v1.kv_layer_groups.parse_kvcache_shape_spec "
+            "for the authoritative parser. Default: '%s'" % _DEFAULT_SHAPE_SPEC
+        ),
+    )
+    kv.add_argument(
+        "--num-blocks",
+        type=int,
+        default=1024,
+        help="Paged blocks (default: 1024)",
+    )
+    kv.add_argument(
+        "--block-size",
+        type=int,
+        default=16,
+        help="Tokens per block (default: 16)",
+    )
+
+    parser.add_argument(
+        "--start",
+        type=int,
+        default=0,
+        help="Starting sequence number (default: 0)",
+    )
+    parser.add_argument(
+        "--end",
+        type=int,
+        default=None,
+        help=("Ending sequence number (exclusive). If not set, runs forever."),
+    )
+    parser.add_argument(
+        "--interval",
+        type=float,
+        default=0.5,
+        help=("Seconds between requests (default: 0.5)"),
+    )
+    parser.add_argument(
+        "--concurrency",
+        type=int,
+        default=1,
+        help=(
+            "Number of clients sending requests at the same time "
+            "(default: 1). Each client opens its own connection, allocates "
+            "its own KV tensors, and registers its Workers under distinct "
+            "instance IDs; sequences are dealt round-robin across clients."
+        ),
+    )
+    parser.add_argument(
+        "--url",
+        default="http://localhost:8080",
+        help=("Legacy HTTP URL; checksums are computed from local worker buffers"),
+    )
+
+    prof = parser.add_argument_group(
+        "server profiling",
+        "Flame-graph the MP server process while this benchmark drives "
+        "load into it. The server's store path (hashing, allocation, "
+        "gather, D2H) runs in its own process, not in this client, so "
+        "profiling attaches to --profile-server-pid rather than to the "
+        "benchmark. See 'lmcache tool flamegraph' for the standalone form.",
+    )
+    prof.add_argument(
+        "--flamegraph",
+        choices=["on", "off"],
+        default="off",
+        help="Record a flame graph of the server during the run (default: off).",
+    )
+    prof.add_argument(
+        "--profile-server-pid",
+        type=int,
+        default=0,
+        metavar="PID",
+        help=(
+            "Server process to profile, e.g. $(pgrep -f 'lmcache server'). "
+            "Required when --flamegraph on."
+        ),
+    )
+    prof.add_argument(
+        "--flamegraph-mode",
+        default="gil",
+        metavar="MODE[,MODE...]",
+        help=(
+            "What to sample in the server (default: gil). Pass several "
+            "comma-separated to profile one load run per mode. Modes: on-cpu, "
+            "off-cpu, wakeup, offwake (perf/bcc), wall, gil (py-spy). perf/bcc "
+            "name Python functions only when the server was launched with "
+            "PYTHONPERFSUPPORT=1. See the 'lmcache tool flamegraph' docs."
+        ),
+    )
+    prof.add_argument(
+        "--flamegraph-output",
+        default="",
+        metavar="PATH",
+        help=(
+            "SVG output path. Default: "
+            "/tmp/lmcache_bench_flames/server-pid<PID>.<mode>.svg."
+        ),
+    )
+    prof.add_argument(
+        "--flamegraph-scripts-dir",
+        default="",
+        metavar="DIR",
+        help=(
+            "Directory with the FlameGraph scripts (flamegraph.pl, "
+            "stackcollapse-perf.pl); default ~/FlameGraph (cloned there on "
+            "first use). Unused by --flamegraph-mode wall / gil, which "
+            "render their own SVG."
+        ),
+    )

lmcache/cli/commands/bench/server_bench/buffers.py

@@ -12,6 +12,14 @@ import torch
 from lmcache import torch_dev, torch_device_type
 from lmcache.v1.kv_layer_groups import KVLayerGroupInfo
 from lmcache.v1.multiprocess.group_view import EngineGroupInfo
+from lmcache.v1.platform import synchronize_device
+
+
+def _block_bytes(tensor: torch.Tensor, dim: int, block_ids: Sequence[int]) -> bytes:
+    """Gather block bytes in logical order."""
+    indices = torch.tensor(block_ids, dtype=torch.long, device=tensor.device)
+    selected = tensor.index_select(dim, indices)
+    return selected.contiguous().view(torch.uint8).cpu().numpy().tobytes()
 
 
 class KVCacheBuffers:
@@ -61,43 +69,51 @@ class KVCacheBuffers:
                     tensor = torch.randint(
                         bounds.min, bounds.max + 1, shape, dtype=group.dtype, device=dev
                     )
-                name = f"layer.{len(self.kv_caches)}"
-                self.kv_caches[name] = tensor
-                self._block_dims[name] = block_dim
+                tensor_name = f"tensor.{len(self.kv_caches)}"
+                self.kv_caches[tensor_name] = tensor
+                self._block_dims[tensor_name] = block_dim
 
-    def block_views(self, block_offset: int, num_blocks: int) -> Iterator[torch.Tensor]:
-        """Yield writable block-range views in layer order."""
-        for name, tensor in self.kv_caches.items():
-            yield tensor.narrow(self._block_dims[name], block_offset, num_blocks)
+    def block_views(
+        self, block_ids_per_group: Sequence[Sequence[int]]
+    ) -> Iterator[torch.Tensor]:
+        """Yield validated block views in group/tensor/logical-block order."""
+        for _, tensor, dim, block_ids in self._iter_group_tensors(block_ids_per_group):
+            for block_id in block_ids:
+                yield tensor.narrow(dim, block_id, 1)
 
-    def zero_blocks(self, block_offset: int, num_blocks: int) -> None:
-        """Zero the selected block range in every layer."""
-        for view in self.block_views(block_offset, num_blocks):
+    def set_blocks_to_zero(self, block_ids_per_group: Sequence[Sequence[int]]) -> None:
+        """Set bound blocks to zero."""
+        for view in self.block_views(block_ids_per_group):
             view.zero_()
 
-    def checksums(
-        self, block_offset: int, num_blocks: int, chunk_size: int
-    ) -> list[str]:
-        """Hash selected CPU blocks per chunk across layers.
+    def checksum(self, block_ids_per_group: Sequence[Sequence[int]]) -> str:
+        """Hash selected blocks across tensors in logical order into one MD5 digest."""
+        digest = hashlib.md5()
+        for _, tensor, dim, block_ids in self._iter_group_tensors(block_ids_per_group):
+            digest.update(_block_bytes(tensor, dim, block_ids))
+        return digest.hexdigest()
 
-        ``chunk_size`` counts tokens; incomplete trailing chunks are ignored.
-        """
-        if chunk_size % self.block_size != 0:
-            raise ValueError(
-                "chunk_size %d must be a multiple of block_size %d"
-                % (chunk_size, self.block_size)
-            )
-        blocks_per_chunk = chunk_size // self.block_size
-        checksums: list[str] = []
-        for chunk in range(num_blocks // blocks_per_chunk):
-            start = block_offset + chunk * blocks_per_chunk
-            digest = hashlib.md5()
-            for view in self.block_views(start, blocks_per_chunk):
-                digest.update(view.contiguous().view(torch.uint8).numpy().tobytes())
-            checksums.append(digest.hexdigest())
-        return checksums
+    def synchronize(self) -> None:
+        """Wait for buffer preparation outside transfer timing."""
+        synchronize_device(next(iter(self.kv_caches.values())).device)
 
     def release(self) -> None:
         """Drop owned references without synchronizing or unregistering."""
         self.kv_caches.clear()
         self._block_dims.clear()
+
+    def _iter_group_tensors(
+        self, block_ids_per_group: Sequence[Sequence[int]]
+    ) -> Iterator[tuple[int, torch.Tensor, int, Sequence[int]]]:
+        """Yield registered tensor indices, tensors, block axes, and bound block IDs."""
+        for group, block_ids in zip(
+            self.engine_group_infos, block_ids_per_group, strict=True
+        ):
+            for tensor_index in group.layer_indices:
+                tensor_name = f"tensor.{tensor_index}"
+                yield (
+                    tensor_index,
+                    self.kv_caches[tensor_name],
+                    self._block_dims[tensor_name],
+                    block_ids,
+                )

lmcache/cli/commands/bench/server_bench/cases/baseline.py

@@ -49,6 +49,7 @@ class BaselineBenchCase:
 
     sequence_count: int | None
     interval_seconds: float
+    num_tokens: int = field(default=512, kw_only=True)
     sequence_id_offset: int = 0
     sequence_id_stride: int = 1
     stop: threading.Event | None = field(default=None, compare=False)
@@ -206,65 +207,71 @@ class BaselineBenchCase:
                     )
                 )
 
-    @staticmethod
     def _execute_request(
+        self,
         client: ServerBenchClient,
         sequence_id: int,
         request_kind: Literal["cold", "warm"],
     ) -> _RequestOutcome | None:
+        # First Party
+        from lmcache.cli.commands.bench.server_bench.helpers import _build_token_ids
+
         request = client.create_request(
             sequence_id,
             request_id="req-%d-%s" % (sequence_id, request_kind),
             request_kind=request_kind,
+            token_ids=_build_token_ids(sequence_id, self.num_tokens),
         )
         if request is None:
             return None
 
-        lookup = client.lookup(request)
-        if not lookup.succeeded:
-            return _RequestOutcome(lookup=lookup)
-
-        hit_tokens = lookup.hit_chunks * request.chunk_size
-        miss_tokens = request.num_full_tokens - hit_tokens
+        try:
+            lookup = client.lookup(request)
+            if not lookup.succeeded:
+                return _RequestOutcome(lookup=lookup)
+
+            hit_tokens = lookup.hit_chunks * request.chunk_size
+            miss_tokens = request.num_full_tokens - hit_tokens
+
+            checksums: dict[int, list[str]] | None = None
+            if request_kind == "cold" and miss_tokens > 0:
+                checksums = client.compute_checksums(
+                    request,
+                    start_token=hit_tokens,
+                    token_count=miss_tokens,
+                    operation="store",
+                )
+            if request_kind == "warm" and hit_tokens > 0:
+                client.zero_destination(
+                    request,
+                    start_token=0,
+                    token_count=hit_tokens,
+                )
 
-        checksums: dict[int, list[str]] | None = None
-        if request_kind == "cold" and miss_tokens > 0:
-            checksums = client.compute_checksums(
-                request,
-                start_token=hit_tokens,
-                token_count=miss_tokens,
-                operation="store",
-            )
-        if request_kind == "warm" and hit_tokens > 0:
-            client.zero_destination(
+            retrieve = client.retrieve(
                 request,
                 start_token=0,
                 token_count=hit_tokens,
             )
-
-        retrieve = client.retrieve(
-            request,
-            start_token=0,
-            token_count=hit_tokens,
-        )
-        store = client.store(
-            request,
-            start_token=hit_tokens,
-            token_count=miss_tokens,
-        )
-
-        if request_kind == "warm" and hit_tokens > 0:
-            checksums = client.compute_checksums(
+            store = client.store(
                 request,
-                start_token=0,
-                token_count=hit_tokens,
-                operation="retrieve",
+                start_token=hit_tokens,
+                token_count=miss_tokens,
             )
 
-        client.end_session(request)
-        return _RequestOutcome(
-            lookup=lookup,
-            checksums=checksums,
-            retrieve=retrieve,
-            store=store,
-        )
+            if request_kind == "warm" and hit_tokens > 0:
+                checksums = client.compute_checksums(
+                    request,
+                    start_token=0,
+                    token_count=hit_tokens,
+                    operation="retrieve",
+                )
+
+            return _RequestOutcome(
+                lookup=lookup,
+                checksums=checksums,
+                retrieve=retrieve,
+                store=store,
+            )
+        finally:
+            client.release_request(request)