wehub-resource-sync
424 lines
17 KiB
Plaintext
424 lines
17 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "1cd06b97",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Setup vLLM and LMCache server\n",
|
|
"Follow below instructions (1-4), then proceed to below Python cell.\n",
|
|
"1. Start LMCache server first\n",
|
|
"```sh\n",
|
|
"lmcache server \\\n",
|
|
" --l1-size-gb 8 \\\n",
|
|
" --eviction-policy LRU \\\n",
|
|
" --chunk-size 256 \\\n",
|
|
" --port 6555 \\\n",
|
|
" --http-port 8080 \\\n",
|
|
" --shm-name lmcache_kvcache_sdk_e2e \\\n",
|
|
" --no-l1-use-lazy\n",
|
|
"```\n",
|
|
"2. Wait until LMCache server is ready\n",
|
|
"```sh\n",
|
|
"curl -sf http://localhost:8080/healthcheck && echo \" LMCache ready\"\n",
|
|
"```\n",
|
|
"3. Start vLLM once LMCache server is ready\n",
|
|
"```sh\n",
|
|
"env -u VLLM_PORT \\\n",
|
|
" CUDA_VISIBLE_DEVICES=0 \\\n",
|
|
" VLLM_ENABLE_V1_MULTIPROCESSING=0 \\\n",
|
|
" VLLM_BATCH_INVARIANT=1 \\\n",
|
|
" PYTHONHASHSEED=0 \\\n",
|
|
" vllm serve Qwen/Qwen3-8B \\\n",
|
|
" --port 8000 \\\n",
|
|
" --served-model-name Qwen/Qwen3-8B \\\n",
|
|
" --no-enable-prefix-caching \\\n",
|
|
" --enforce-eager \\\n",
|
|
" --max-model-len 4096 \\\n",
|
|
" --gpu-memory-utilization 0.6 \\\n",
|
|
" --kv-transfer-config '{\"kv_connector\":\"LMCacheMPConnector\",\"kv_role\":\"kv_both\",\"kv_load_failure_policy\":\"recompute\",\"kv_connector_extra_config\":{\"lmcache.mp.host\":\"tcp://localhost\",\"lmcache.mp.port\":6555,\"lmcache.mp.mq_timeout\":10}}' \\\n",
|
|
" --override-generation-config '{\"temperature\": 0}' \\\n",
|
|
" --trust-remote-code\n",
|
|
"```\n",
|
|
"4. Wait until vLLM is ready\n",
|
|
"```sh\n",
|
|
"curl -sf http://localhost:8000/v1/models && echo \" vLLM ready\"\n",
|
|
"```"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "c2d81202",
|
|
"metadata": {},
|
|
"source": [
|
|
"## Run E2E KV Edit Example"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 1,
|
|
"id": "aa87a653",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"/home/rani/LMCache/.venv/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
|
|
" from .autonotebook import tqdm as notebook_tqdm\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:41,202] LMCache INFO:\u001b[0m torch_dev=<module 'torch.cuda' from '/home/rani/LMCache/.venv/lib/python3.12/site-packages/torch/cuda/__init__.py'>, torch_device_type=cuda \u001b[3m(__init__.py:62:lmcache)\u001b[0m\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:41,267] LMCache INFO:\u001b[0m Skipping backend lmcache.v1.platform.musa.ops: predicate returned False \u001b[3m(__init__.py:102:lmcache)\u001b[0m\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:41,268] LMCache INFO:\u001b[0m Skipping backend lmcache.xpu_ops: predicate returned False \u001b[3m(__init__.py:102:lmcache)\u001b[0m\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:41,269] LMCache INFO:\u001b[0m Using backend: lmcache.c_ops \u001b[3m(__init__.py:120:lmcache)\u001b[0m\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:41,604] LMCache INFO:\u001b[0m multi_layer_block_kv_transfer mode: ptr \u001b[3m(base.py:94:lmcache.v1.multiprocess.transfer_context.base)\u001b[0m\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# SPDX-License-Identifier: Apache-2.0\n",
|
|
"\"\"\"End-to-end KV cache remapping driver for the SDK example.\"\"\"\n",
|
|
"\n",
|
|
"# Standard\n",
|
|
"from dataclasses import dataclass\n",
|
|
"from itertools import islice\n",
|
|
"import json\n",
|
|
"import time\n",
|
|
"from typing import cast\n",
|
|
"\n",
|
|
"# Third Party\n",
|
|
"import httpx\n",
|
|
"from transformers import AutoTokenizer, PreTrainedTokenizerBase\n",
|
|
"\n",
|
|
"# First Party\n",
|
|
"import lmcache.sdk.kvcache as lmc_sdk"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 2,
|
|
"id": "8f1065b0",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"model_name = \"Qwen/Qwen3-8B\"\n",
|
|
"vllm_url = \"http://localhost:8000\"\n",
|
|
"lmcache_url = \"http://localhost:8080\"\n",
|
|
"lmcache_mq_url = \"tcp://localhost:6555\"\n",
|
|
"chunk_size = 256\n",
|
|
"min_prompt_tokens = chunk_size * 2\n",
|
|
"fake_prefix_tokens = 32\n",
|
|
"max_tokens = 32\n",
|
|
"timeout = 60 # seconds\n",
|
|
"trust_remote_code = True"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 3,
|
|
"id": "f894f55b",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"@dataclass(frozen=True)\n",
|
|
"class CompletionResult:\n",
|
|
" \"\"\"Text and latency returned by one OpenAI-compatible completion call.\"\"\"\n",
|
|
"\n",
|
|
" text: str\n",
|
|
" elapsed_seconds: float\n",
|
|
"\n",
|
|
"\n",
|
|
"def _post_completion(\n",
|
|
" *,\n",
|
|
" vllm_url: str,\n",
|
|
" model_name: str,\n",
|
|
" prompt: str | list[int],\n",
|
|
" max_tokens: int,\n",
|
|
" timeout: float,\n",
|
|
") -> CompletionResult:\n",
|
|
" \"\"\"Send one non-streaming completion request to vLLM.\"\"\"\n",
|
|
" payload = {\n",
|
|
" \"model\": model_name,\n",
|
|
" \"prompt\": prompt,\n",
|
|
" \"max_tokens\": max_tokens,\n",
|
|
" \"min_tokens\": max_tokens,\n",
|
|
" \"temperature\": 0,\n",
|
|
" \"seed\": 0,\n",
|
|
" \"ignore_eos\": True,\n",
|
|
" }\n",
|
|
" start = time.perf_counter()\n",
|
|
" response = httpx.post(\n",
|
|
" f\"{vllm_url.rstrip('/')}/v1/completions\",\n",
|
|
" json=payload,\n",
|
|
" timeout=timeout,\n",
|
|
" )\n",
|
|
" elapsed = time.perf_counter() - start\n",
|
|
" response.raise_for_status()\n",
|
|
" body = response.json()\n",
|
|
" choices = body.get(\"choices\")\n",
|
|
" if not isinstance(choices, list) or not choices:\n",
|
|
" raise RuntimeError(f\"completion response missing choices: {body}\")\n",
|
|
" first_choice = choices[0]\n",
|
|
" if not isinstance(first_choice, dict) or not isinstance(\n",
|
|
" first_choice.get(\"text\"), str\n",
|
|
" ):\n",
|
|
" raise RuntimeError(f\"completion response has invalid choice: {body}\")\n",
|
|
" return CompletionResult(text=first_choice[\"text\"], elapsed_seconds=elapsed)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "4cd588ce",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"SOURCE_PARAGRAPH = (\n",
|
|
" \"A systems researcher is studying how an inference cache changes the \"\n",
|
|
" \"latency profile of a long language model prompt. The notes discuss \"\n",
|
|
" \"attention keys, attention values, memory tiers, token chunks, and the \"\n",
|
|
" \"careful measurement of cold and warm requests.\"\n",
|
|
")\n",
|
|
"\n",
|
|
"\n",
|
|
"def _build_prompts(\n",
|
|
" tokenizer: PreTrainedTokenizerBase,\n",
|
|
" *,\n",
|
|
" min_prompt_tokens: int,\n",
|
|
" chunk_size: int,\n",
|
|
" fake_prefix_tokens: int,\n",
|
|
") -> tuple[list[int], list[int], int]:\n",
|
|
" \"\"\"Build equal-length source/target prompts and the cached-prefix length.\n",
|
|
" The prompts differ only in ``fake_prefix_tokens`` synthetic leading token IDs.\"\"\"\n",
|
|
" if not 0 < fake_prefix_tokens < chunk_size:\n",
|
|
" raise ValueError(\n",
|
|
" f\"fake_prefix_tokens must be in (0, {chunk_size}), got {fake_prefix_tokens}\"\n",
|
|
" )\n",
|
|
" min_cache = max(1, min_prompt_tokens - fake_prefix_tokens)\n",
|
|
" cache_tokens = ((min_cache + chunk_size - 1) // chunk_size) * chunk_size\n",
|
|
"\n",
|
|
" # Two distinct, non-special lead IDs for the source/target prefixes.\n",
|
|
" special = {int(t) for t in tokenizer.all_special_ids}\n",
|
|
" candidates = (t for t in range(1000, int(tokenizer.vocab_size)) if t not in special)\n",
|
|
" leads = list(islice(candidates, 2))\n",
|
|
" if len(leads) < 2:\n",
|
|
" raise ValueError(\"could not find two usable non-special token IDs\")\n",
|
|
" source_lead, target_lead = leads\n",
|
|
"\n",
|
|
" # Repeat the paragraph until it covers the suffix, then take the trailing slice.\n",
|
|
" suffix_len = cache_tokens - fake_prefix_tokens\n",
|
|
" text = SOURCE_PARAGRAPH\n",
|
|
" suffix = tokenizer.encode(text, add_special_tokens=False)\n",
|
|
" while len(suffix) < suffix_len:\n",
|
|
" text = f\"{text}\\n\\n{SOURCE_PARAGRAPH}\"\n",
|
|
" suffix = tokenizer.encode(text, add_special_tokens=False)\n",
|
|
" suffix = [int(t) for t in suffix[-suffix_len:]]\n",
|
|
"\n",
|
|
" source = [source_lead] * fake_prefix_tokens + suffix\n",
|
|
" target = [target_lead] * fake_prefix_tokens + suffix\n",
|
|
" return source, target, cache_tokens"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 5,
|
|
"id": "030145fc",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stderr",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"\u001b[32;20m[2026-06-23 23:26:52,200] LMCache INFO:\u001b[0m Initialized LMCacheKVCacheContext with instance_id=3403760, model_name=Qwen/Qwen3-8B, chunk_size=256, shm_name=lmcache_l1_pool_lmcache_kvcache_sdk_e2e \u001b[3m(kvcache.py:104:lmcache.sdk.kvcache)\u001b[0m\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:52,206] LMCache INFO:\u001b[0m Creating transfer context (device_type=cpu, mode=auto) \u001b[3m(worker_transfer.py:551:lmcache.v1.multiprocess.transfer_context.worker_transfer)\u001b[0m\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:52,206] LMCache INFO:\u001b[0m Engine KV Format: EngineKVFormat.NL_X_NB_TWO_NH_BS_HS NL x [NB, 2, NH, BS, HS] \u001b[3m(detection.py:44:lmcache.v1.gpu_connector.kv_format.detection)\u001b[0m\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:52,208] LMCache INFO:\u001b[0m Creating EngineDrivenContextShm (shm_name=lmcache_l1_pool_lmcache_kvcache_sdk_e2e, pool_size=150323855360) \u001b[3m(base.py:235:lmcache.v1.multiprocess.transfer_context.base)\u001b[0m\n",
|
|
"\u001b[32;20m[2026-06-23 23:26:52,209] LMCache INFO:\u001b[0m Worker non-GPU transfer context registered (instance_id=3403760, mode=SHM) \u001b[3m(worker_transfer.py:420:lmcache.v1.multiprocess.transfer_context.worker_transfer)\u001b[0m\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"# Only create one LMCacheKVCacheContext, reuse for retrieve() and store()\n",
|
|
"# multiple times. Only close() once at the end.\n",
|
|
"ctx = lmc_sdk.connect(\n",
|
|
" url=lmcache_mq_url,\n",
|
|
" http_url=lmcache_url,\n",
|
|
" model_name=model_name,\n",
|
|
" timeout=timeout,\n",
|
|
")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 6,
|
|
"id": "2f1c8d9a",
|
|
"metadata": {},
|
|
"outputs": [
|
|
{
|
|
"name": "stdout",
|
|
"output_type": "stream",
|
|
"text": [
|
|
"== Step 1: source inference stores KV under source token IDs ==\n",
|
|
"== Step 2: retrieve source KV into memory through lmcache.sdk ==\n",
|
|
"== Step 3: store source KV under different target token IDs ==\n",
|
|
"== Step 4: target inference reuses the remapped token IDs ==\n",
|
|
"== Evaluation ==\n",
|
|
"{\n",
|
|
" \"vllm_model_name\": \"Qwen/Qwen3-8B\",\n",
|
|
" \"cache_tokens\": 512,\n",
|
|
" \"source_prompt_tokens\": 512,\n",
|
|
" \"target_prompt_tokens\": 512,\n",
|
|
" \"retrieved_hit_tokens\": 512,\n",
|
|
" \"retrieved_hit_chunks\": 2,\n",
|
|
" \"source_target_same_length\": true,\n",
|
|
" \"source_target_last_hit_tokens_match\": false,\n",
|
|
" \"target_prefix_ids_differ_from_source\": true,\n",
|
|
" \"outputs_match\": true,\n",
|
|
" \"store_result\": true,\n",
|
|
" \"source_retrieve_after_inference\": {\n",
|
|
" \"shape\": [\n",
|
|
" 2,\n",
|
|
" 36,\n",
|
|
" 512,\n",
|
|
" 1024\n",
|
|
" ],\n",
|
|
" \"dtype\": \"torch.bfloat16\",\n",
|
|
" \"hit_tokens\": 512,\n",
|
|
" \"hit_chunks\": 2\n",
|
|
" },\n",
|
|
" \"source_latency_seconds\": 1.2143411422148347,\n",
|
|
" \"target_latency_seconds\": 1.1793114291504025,\n",
|
|
" \"source_output_preview\": \"Okay, I need to understand what the user is asking for. They provided a series of repeated sentences about a systems researcher studying how an inference cache affects the\",\n",
|
|
" \"target_output_preview\": \"Okay, I need to understand what the user is asking for. They provided a series of repeated sentences about a systems researcher studying how an inference cache affects the\"\n",
|
|
"}\n"
|
|
]
|
|
}
|
|
],
|
|
"source": [
|
|
"tokenizer = AutoTokenizer.from_pretrained(\n",
|
|
" model_name,\n",
|
|
" trust_remote_code=trust_remote_code,\n",
|
|
")\n",
|
|
"source_tokens, target_tokens, cache_tokens = _build_prompts(\n",
|
|
" tokenizer,\n",
|
|
" min_prompt_tokens=min_prompt_tokens,\n",
|
|
" chunk_size=chunk_size,\n",
|
|
" fake_prefix_tokens=fake_prefix_tokens,\n",
|
|
")\n",
|
|
"\n",
|
|
"print(\"== Step 1: source inference stores KV under source token IDs ==\")\n",
|
|
"source_completion = _post_completion(\n",
|
|
" vllm_url=vllm_url,\n",
|
|
" model_name=model_name,\n",
|
|
" prompt=source_tokens,\n",
|
|
" max_tokens=max_tokens,\n",
|
|
" timeout=timeout,\n",
|
|
")\n",
|
|
"\n",
|
|
"print(\"== Step 2: retrieve source KV into memory through lmcache.sdk ==\")\n",
|
|
"retrieved_kv = lmc_sdk.retrieve(\n",
|
|
" ctx=ctx,\n",
|
|
" tokens=source_tokens,\n",
|
|
")\n",
|
|
"if retrieved_kv is None:\n",
|
|
" raise RuntimeError(\"source retrieve missed the expected cached prefix\")\n",
|
|
"retrieved_hit_tokens = int(retrieved_kv.shape[2])\n",
|
|
"retrieved_hit_chunks = retrieved_hit_tokens // chunk_size\n",
|
|
"if retrieved_hit_tokens < cache_tokens:\n",
|
|
" raise RuntimeError(\n",
|
|
" \"source retrieve did not return the expected cached prefix: \"\n",
|
|
" f\"{retrieved_hit_tokens} < {cache_tokens}\"\n",
|
|
" )\n",
|
|
"\n",
|
|
"target_prefix_tokens = target_tokens[:retrieved_hit_tokens]\n",
|
|
"source_prefix_tokens = source_tokens[:retrieved_hit_tokens]\n",
|
|
"if source_prefix_tokens == target_prefix_tokens:\n",
|
|
" raise RuntimeError(\"source and target token prefixes unexpectedly match\")\n",
|
|
"\n",
|
|
"print(\"== Step 3: store source KV under different target token IDs ==\")\n",
|
|
"store_result = lmc_sdk.store(\n",
|
|
" ctx=ctx,\n",
|
|
" kv=retrieved_kv,\n",
|
|
" tokens=target_prefix_tokens,\n",
|
|
")\n",
|
|
"if not store_result:\n",
|
|
" raise RuntimeError(\n",
|
|
" f\"Failed to store {len(target_prefix_tokens)} target prefix tokens. \"\n",
|
|
" f\"Might be because LMCache already has the KV cache.\"\n",
|
|
" )\n",
|
|
"\n",
|
|
"print(\"== Step 4: target inference reuses the remapped token IDs ==\")\n",
|
|
"target_completion = _post_completion(\n",
|
|
" vllm_url=vllm_url,\n",
|
|
" model_name=model_name,\n",
|
|
" prompt=target_tokens,\n",
|
|
" max_tokens=max_tokens,\n",
|
|
" timeout=timeout,\n",
|
|
")\n",
|
|
"outputs_match = source_completion.text == target_completion.text\n",
|
|
"\n",
|
|
"evaluation = {\n",
|
|
" \"vllm_model_name\": model_name,\n",
|
|
" \"cache_tokens\": cache_tokens,\n",
|
|
" \"source_prompt_tokens\": len(source_tokens),\n",
|
|
" \"target_prompt_tokens\": len(target_tokens),\n",
|
|
" \"retrieved_hit_tokens\": retrieved_hit_tokens,\n",
|
|
" \"retrieved_hit_chunks\": retrieved_hit_chunks,\n",
|
|
" \"source_target_same_length\": len(source_tokens) == len(target_tokens),\n",
|
|
" \"source_target_last_hit_tokens_match\": source_tokens[-retrieved_hit_tokens:]\n",
|
|
" == target_tokens[-retrieved_hit_tokens:],\n",
|
|
" \"target_prefix_ids_differ_from_source\": source_prefix_tokens\n",
|
|
" != target_prefix_tokens,\n",
|
|
" \"outputs_match\": outputs_match,\n",
|
|
" \"store_result\": store_result,\n",
|
|
" \"source_retrieve_after_inference\": {\n",
|
|
" \"shape\": tuple(retrieved_kv.shape),\n",
|
|
" \"dtype\": str(retrieved_kv.dtype),\n",
|
|
" \"hit_tokens\": retrieved_hit_tokens,\n",
|
|
" \"hit_chunks\": retrieved_hit_chunks,\n",
|
|
" },\n",
|
|
" \"source_latency_seconds\": source_completion.elapsed_seconds,\n",
|
|
" \"target_latency_seconds\": target_completion.elapsed_seconds,\n",
|
|
" \"source_output_preview\": \" \".join(source_completion.text.split()[:max_tokens]),\n",
|
|
" \"target_output_preview\": \" \".join(target_completion.text.split()[:max_tokens]),\n",
|
|
"}\n",
|
|
"print(\"== Evaluation ==\")\n",
|
|
"print(json.dumps(cast(dict[str, object], evaluation), indent=2, default=str))\n",
|
|
"if not outputs_match:\n",
|
|
" raise RuntimeError(\"Target output did not match source output after KV remap.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": 7,
|
|
"id": "f83db925",
|
|
"metadata": {},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Only close the context when done.\n",
|
|
"lmc_sdk.close(ctx)"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"kernelspec": {
|
|
"display_name": "LMCache (3.12.3.final.0)",
|
|
"language": "python",
|
|
"name": "python3"
|
|
},
|
|
"language_info": {
|
|
"codemirror_mode": {
|
|
"name": "ipython",
|
|
"version": 3
|
|
},
|
|
"file_extension": ".py",
|
|
"mimetype": "text/x-python",
|
|
"name": "python",
|
|
"nbconvert_exporter": "python",
|
|
"pygments_lexer": "ipython3",
|
|
"version": "3.12.3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 5
|
|
}
|