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

lmcache/__init__.py

@@ -15,7 +15,43 @@ except ImportError:
     __version__ = "unknown"
 
 logger = init_logger(__name__)
-__all__ = ["__version__"]
+# Standard
+
+__all__ = ["__version__", "torch_dev", "torch_device_type"]
+
+
+# --------------------------
+# Device detection
+# --------------------------
+def _detect_device() -> tuple[Any, str]:
+    """
+    Detect the available accelerator and return the corresponding torch
+    device module and device type string.
+
+    Returns:
+        tuple[Any, str]: A tuple of (torch_device_module, device_type_string),
+            e.g. ``(torch.cuda, "cuda")`` or ``(torch.xpu, "xpu")``.
+
+    Raises:
+        RuntimeError: If no supported accelerator is found (checked CUDA, XPU, HPU).
+    """
+    try:
+        # Third Party
+        import torch
+    except ImportError:
+        return None, "cpu"  # fallback，CLI-only
+
+    if hasattr(torch, "xpu") and torch.xpu.is_available():
+        return torch.xpu, "xpu"
+    elif hasattr(torch, "hpu") and torch.hpu.is_available():
+        return torch.hpu, "hpu"
+    else:
+        # Fallback: always return torch.cuda for backward compatibility
+        # with existing tests and code paths that assume CUDA is the default.
+        return torch.cuda, "cuda"
+
+
+torch_dev, torch_device_type = _detect_device()
 
 
 # --------------------------

lmcache/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.4.5.dev0'
-__version_tuple__ = version_tuple = (0, 4, 5, 'dev0')
+__version__ = version = '0.4.6.dev0'
+__version_tuple__ = version_tuple = (0, 4, 6, 'dev0')
 
-__commit_id__ = commit_id = 'g6fbec463e'
+__commit_id__ = commit_id = 'g168bbb976'

lmcache/cli/commands/bench/__init__.py

@@ -28,12 +28,28 @@ from lmcache.cli.commands.bench.engine_bench.stats import (
 from lmcache.cli.commands.bench.engine_bench.workloads import create_workload
 from lmcache.logging import init_logger
 
+# Gated for slim install — TestCacheCommand pulls torch/MP runtime.
+_TEST_CACHE_IMPORT_ERROR: ImportError | None = None
+try:
+    # First Party
+    from lmcache.cli.commands.test_cache import TestCacheCommand
+except ImportError as _exc:
+    _TEST_CACHE_IMPORT_ERROR = _exc
+    TestCacheCommand = None  # type: ignore[assignment,misc]
+
 logger = init_logger(__name__)
 
 
 class BenchCommand(BaseCommand):
     """CLI command for sustained performance benchmarking."""
 
+    def __init__(self) -> None:
+        super().__init__()
+        # None on slim install; _register_kvcache registers a stub instead.
+        self._kvcache_delegate = (
+            TestCacheCommand() if TestCacheCommand is not None else None
+        )
+
     def name(self) -> str:
         return "bench"
 
@@ -52,9 +68,10 @@ class BenchCommand(BaseCommand):
         inner = parser.add_subparsers(
             dest="bench_target",
             required=True,
-            metavar="{engine}",
+            metavar="{engine,kvcache}",
         )
         self._register_engine(inner)
+        self._register_kvcache(inner)
 
     def _register_engine(
         self,
@@ -99,6 +116,7 @@ class BenchCommand(BaseCommand):
                 "long-doc-permutator",
                 "long-doc-qa",
                 "multi-round-chat",
+                "prefix-suffix-tuner",
                 "random-prefill",
             ],
             help="Workload type.",
@@ -262,6 +280,36 @@ class BenchCommand(BaseCommand):
             help="Benchmark duration in seconds (default: 60).",
         )
 
+        # --- Prefix-suffix-tuner workload args ---
+        psf_group = parser.add_argument_group(
+            "prefix-suffix-tuner workload options",
+        )
+        psf_group.add_argument(
+            "--psf-context-length",
+            type=int,
+            default=8000,
+            help="Total tokens per request (prefix + breaker + suffix) "
+            "(default: 8000).",
+        )
+        psf_group.add_argument(
+            "--psf-prefix-ratio",
+            type=float,
+            default=0.8,
+            help="Fraction of context-length used by the prefix (default: 0.8). "
+            "Must be in (0.0, 1.0). The remainder (minus a 32-token breaker) is "
+            "the shared suffix.",
+        )
+        psf_group.add_argument(
+            "--psf-thrash",
+            type=float,
+            default=20.0,
+            help="Size in GB of the KV-cache tier to overflow (default: 20.0). "
+            "The workload sizes its prefix pool to slightly more than this, "
+            "so every pass-2 request misses that tier and falls through to "
+            "the next one. Use the L0 (HBM) size for vanilla vLLM baselines, "
+            "or the L1 (LMCache DRAM) size for tiered baselines.",
+        )
+
         # --- Random-prefill workload args ---
         rp_group = parser.add_argument_group(
             "random-prefill workload options",
@@ -281,8 +329,59 @@ class BenchCommand(BaseCommand):
 
         parser.set_defaults(func=self.execute)
 
+    # ------------------------------------------------------------------
+    # kvcache bench target — end-to-end MP cache sanity test
+    # ------------------------------------------------------------------
+
+    def _register_kvcache(
+        self,
+        subparsers: argparse._SubParsersAction,
+    ) -> None:
+        """Register ``lmcache bench kvcache``. Delegates to
+        :class:`TestCacheCommand`, or registers a stub on slim install.
+        """
+        if self._kvcache_delegate is None:
+            subparsers.add_parser(
+                "kvcache",
+                help="(requires full lmcache install)",
+                description=(
+                    "End-to-end sanity test for the LMCache MP cache server. "
+                    "Requires the full `lmcache` package; not available in "
+                    "the `lmcache-cli` install."
+                ),
+            ).set_defaults(func=self.execute)
+            return
+        parser = subparsers.add_parser(
+            "kvcache",
+            help=self._kvcache_delegate.help(),
+            description=(
+                "End-to-end sanity test for the LMCache MP cache server: "
+                "runs LOOKUP / STORE / RETRIEVE against a live MP server "
+                "and verifies KV cache checksums."
+            ),
+        )
+        self._kvcache_delegate.add_arguments(parser)
+        parser.set_defaults(func=self.execute)
+
+    def _bench_kvcache(self, args: argparse.Namespace) -> None:
+        """Dispatch ``lmcache bench kvcache`` to ``TestCacheCommand``."""
+        if self._kvcache_delegate is None:
+            print(
+                "ERROR: `lmcache bench kvcache` needs the full LMCache "
+                "package, but only the `lmcache-cli` shell is installed.\n"
+                "  Install the full package with `pip install lmcache` "
+                "and try again.\n"
+                f"  Original import error: {_TEST_CACHE_IMPORT_ERROR}",
+                file=sys.stderr,
+            )
+            sys.exit(1)
+        self._kvcache_delegate.execute(args)
+
     def execute(self, args: argparse.Namespace) -> None:
-        handlers = {"engine": self._bench_engine}
+        handlers = {
+            "engine": self._bench_engine,
+            "kvcache": self._bench_kvcache,
+        }
         handler = handlers.get(args.bench_target)
         if handler is None:
             print(

lmcache/cli/commands/bench/engine_bench/config.py

@@ -93,7 +93,7 @@ def auto_detect_model(engine_url: str) -> str:
 
 
 def _fetch_lmcache_status(lmcache_url: str) -> dict:
-    """Fetch ``/api/status`` from the LMCache HTTP server.
+    """Fetch ``/status`` from the LMCache HTTP server.
 
     Returns:
         Parsed JSON response.
@@ -104,7 +104,7 @@ def _fetch_lmcache_status(lmcache_url: str) -> dict:
     url = lmcache_url.rstrip("/")
     if not url.startswith(("http://", "https://")):
         url = f"http://{url}"
-    status_url = f"{url}/api/status"
+    status_url = f"{url}/status"
 
     logger.debug("Fetching LMCache status from %s", status_url)
 
@@ -125,7 +125,7 @@ def _find_model_meta(
     """Find the GPU metadata entry matching *model_name*.
 
     Args:
-        gpu_meta: The ``gpu_context_meta`` dict from ``/api/status``.
+        gpu_meta: The ``gpu_context_meta`` dict from ``/status``.
         model_name: Model name to match.
 
     Returns:
@@ -148,7 +148,7 @@ def _find_model_meta(
 def resolve_tokens_per_gb(lmcache_url: str, model_name: str) -> int:
     """Query the LMCache server and compute tokens per GB of KV cache.
 
-    Fetches ``/api/status``, finds the model entry matching
+    Fetches ``/status``, finds the model entry matching
     *model_name*, and computes::
 
         global_bytes_per_token = cache_size_per_token * world_size
@@ -172,6 +172,10 @@ def resolve_tokens_per_gb(lmcache_url: str, model_name: str) -> int:
     data = _fetch_lmcache_status(lmcache_url)
 
     gpu_meta = data.get("gpu_context_meta", {})
+    if not gpu_meta:
+        # CB-only deployments (engine_type="blend") populate
+        # cb_gpu_context_meta instead of gpu_context_meta.
+        gpu_meta = data.get("cb_gpu_context_meta", {})
     if not gpu_meta:
         raise RuntimeError(
             "No model info returned by LMCache server; "

lmcache/cli/commands/bench/engine_bench/interactive/schema.py

@@ -112,6 +112,10 @@ ALL_ITEMS: list[ConfigItem] = [
             ),
             ("long-doc-qa", "Repeated Q&A over long documents (tests KV cache reuse)"),
             ("multi-round-chat", "Multi-turn chat with stateful sessions"),
+            (
+                "prefix-suffix-tuner",
+                "Two-pass sequential workload demonstrating tiered KV cache reuse",
+            ),
             ("random-prefill", "Prefill-only requests fired simultaneously"),
         ],
         phase=PHASE_REQUIRED,
@@ -317,6 +321,43 @@ ALL_ITEMS: list[ConfigItem] = [
         condition=_workload_is("multi-round-chat"),
         phase=PHASE_WORKLOAD,
     ),
+    # ── Phase 3: prefix-suffix-tuner ──────────────────────────────────
+    ConfigItem(
+        key="psf_context_length",
+        display_name="Context length (tokens)",
+        description="Total tokens per request (prefix + breaker + suffix).",
+        input_type="int",
+        default=8000,
+        condition=_workload_is("prefix-suffix-tuner"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="psf_prefix_ratio",
+        display_name="Prefix ratio",
+        description=(
+            "Fraction of context-length used by the prefix. Must be in "
+            "(0.0, 1.0). The remainder (minus a 32-token breaker) is the "
+            "shared suffix."
+        ),
+        input_type="float",
+        default=0.8,
+        condition=_workload_is("prefix-suffix-tuner"),
+        phase=PHASE_WORKLOAD,
+    ),
+    ConfigItem(
+        key="psf_thrash",
+        display_name="Target tier size (GB)",
+        description=(
+            "Size in GB of the KV-cache tier to overflow. The prefix pool "
+            "is sized to slightly more than this, so every pass-2 request "
+            "misses the targeted tier. Use the L0 (HBM) size for vanilla "
+            "vLLM, or the L1 (LMCache DRAM) size for tiered baselines."
+        ),
+        input_type="float",
+        default=20.0,
+        condition=_workload_is("prefix-suffix-tuner"),
+        phase=PHASE_WORKLOAD,
+    ),
     # ── Phase 3: random-prefill ───────────────────────────────────────
     ConfigItem(
         key="rp_request_length",

lmcache/cli/commands/bench/engine_bench/request_sender.py

@@ -111,6 +111,7 @@ class RequestSender:
         """
         submit_time = time.time()
         first_token_time = 0.0
+        first_chunk_time = 0.0
         tokens: list[str] = []
         num_input_tokens = 0
         num_output_tokens = 0
@@ -119,6 +120,9 @@ class RequestSender:
             response = await self._create_stream(messages, max_tokens)
 
             async for chunk in response:
+                if not first_chunk_time:
+                    first_chunk_time = time.time()
+
                 # Extract usage from final chunk
                 usage = getattr(chunk, "usage", None)
                 if usage is not None:
@@ -136,6 +140,11 @@ class RequestSender:
                     tokens.append(content)
 
             finish_time = time.time()
+            if first_token_time == 0.0 and num_output_tokens > 0:
+                # Empty-content stream (common with max_tokens=1, e.g. the
+                # single token is EOS): use first chunk arrival as TTFT —
+                # closer to engine prefill completion than finish_time.
+                first_token_time = first_chunk_time or finish_time
             successful = first_token_time > 0.0
             ttft = (first_token_time - submit_time) if successful else -1.0
             request_latency = finish_time - submit_time

lmcache/cli/commands/bench/engine_bench/workloads/__init__.py

@@ -30,6 +30,10 @@ from lmcache.cli.commands.bench.engine_bench.workloads.multi_round_chat import (
     MultiRoundChatConfig,
     MultiRoundChatWorkload,
 )
+from lmcache.cli.commands.bench.engine_bench.workloads.prefix_suffix_tuner import (
+    PrefixSuffixTunerConfig,
+    PrefixSuffixTunerWorkload,
+)
 from lmcache.cli.commands.bench.engine_bench.workloads.random_prefill import (
     RandomPrefillConfig,
     RandomPrefillWorkload,
@@ -43,6 +47,8 @@ __all__ = [
     "LongDocQAWorkload",
     "MultiRoundChatConfig",
     "MultiRoundChatWorkload",
+    "PrefixSuffixTunerConfig",
+    "PrefixSuffixTunerWorkload",
     "RandomPrefillConfig",
     "RandomPrefillWorkload",
     "create_workload",
@@ -52,6 +58,7 @@ _WORKLOAD_NAMES = (
     "long-doc-permutator",
     "long-doc-qa",
     "multi-round-chat",
+    "prefix-suffix-tuner",
     "random-prefill",
 )
 
@@ -135,6 +142,22 @@ def create_workload(
             seed=config.seed,
         )
 
+    if config.workload == "prefix-suffix-tuner":
+        psf_workload_config = PrefixSuffixTunerConfig.resolve(
+            tokens_per_gb_kvcache=config.tokens_per_gb_kvcache,
+            context_length=args.psf_context_length,
+            prefix_ratio=args.psf_prefix_ratio,
+            thrash=args.psf_thrash,
+        )
+        return PrefixSuffixTunerWorkload(
+            config=psf_workload_config,
+            request_sender=request_sender,
+            stats_collector=stats_collector,
+            progress_monitor=progress_monitor,
+            seed=config.seed,
+            model_name=config.model,
+        )
+
     if config.workload == "random-prefill":
         rp_workload_config = RandomPrefillConfig.resolve(
             request_length=args.rp_request_length,