Skip to content

docs(kv-caching): Run every step one replica at a time #96

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
alexey-temnikov wants to merge 1 commit into from
Closed

docs(kv-caching): Run every step one replica at a time #96

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
72 changes: 50 additions & 22 deletions cookbooks/kv-caching/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -51,8 +51,9 @@ to the other.

- **A container runtime with Compose** (Docker Compose or compatible).
- **Memory.** Give your container runtime about **9 GiB** for the default
two-replica walkthrough, or about **5.5 GiB** to run a single replica (see
[Running on a smaller machine](#running-on-a-smaller-machine-one-replica)).
two-replica walkthrough, or about **5.5 GiB** to run every step with one
replica at a time (see
[Running on a smaller machine](#running-on-a-smaller-machine)).
See [Memory requirements](#memory-requirements) for the measured numbers.
- **Python 3.12** for the notebook client.
- **Disk.** The first run downloads about 3.5 GiB on amd64 (2.5 GiB on arm64):
Expand All @@ -69,10 +70,10 @@ The lowest memory limit at which every cold run of the full walkthrough passed
the containers under one hard memory limit and no swap. Docker, Podman and
nerdctl measured the same within 0.25 GiB:

| CPU | two replicas | one replica |
| --- | --- | --- |
| x86-64 | 6.75 GiB | 3.75 GiB |
| arm64 (AWS Graviton) | 7.5 GiB | 4.5 GiB |
| CPU | two replicas | one replica at a time | single replica |
| --- | --- | --- | --- |
| x86-64 | 6.75 GiB | 3.75 GiB | 3.75 GiB |
| arm64 (AWS Graviton) | 7.5 GiB | 4.5 GiB | 4.5 GiB |

Docker Desktop and Podman machine run the containers in a VM, which needs memory
of its own on top of these numbers, and Apple silicon was not measured
Expand Down Expand Up @@ -154,24 +155,50 @@ watch the cache fill and get reused.
docker compose down
```

## Running on a smaller machine (one replica)
## Running on a smaller machine

The default walkthrough runs two vLLM replicas to show cross-instance cache
sharing, which needs the most memory (see [Prerequisites](#prerequisites)). If
your machine has less, run a single replica instead. You still see the core
win — a cold turn stores its KV in Valkey and a warm turn reads it back — just
not cross-replica sharing.
The default walkthrough keeps both vLLM replicas running at the same time,
which needs the most memory (see [Prerequisites](#prerequisites)). If your
container runtime has less, run the replicas one at a time: only one replica is
ever resident, so it fits the one-replica figure in Prerequisites.

Start only replica A (Compose also starts `valkey` and `lmcache-server`, which
it depends on):
### One replica at a time (every step, including cross-replica reuse)

```bash
docker compose up -d --wait vllm-a
```
Valkey and the LMCache server keep running the whole time, so the KV that
replica A wrote is still in Valkey when replica B starts. Replica B reading it
back also shows that the cache outlives the replica that computed it.

1. Start Valkey, the LMCache server, and replica A only:

```bash
docker compose up -d --wait vllm-a
```

2. Run the notebook from the top and stop at
**"A second replica hits the first replica's cache."**
3. Swap replica A for replica B, from a terminal in this directory:

```bash
docker compose stop vllm-a
docker compose up -d --wait --no-deps vllm-b
```

Keep `--no-deps`: `vllm-b` normally waits for `vllm-a` to be healthy, so
without it Compose would start replica A again. Replica B needs a minute or
two to install LMCache, load the model, and warm up.
4. Run the remaining cells. Replica B's request reads replica A's KV from
Valkey: the `GET` calls go up and `dbsize` stays the same.

With Podman, run the same commands with `podman compose`. If your
`podman compose` has no `--wait`, drop it and wait until `podman compose ps`
shows `vllm-b` as healthy.

### A single replica

Then run the notebook as usual, but skip the one cell that calls the second
replica — it's the cell under **"A second replica hits the first replica's
cache,"** marked in the notebook as the only cell to skip in single-replica mode.
If you only want the single-replica part of the demo, start replica A as in
step 1 and skip the one cell that calls `replica_b` (under
**"A second replica hits the first replica's cache"**). You still see a cold
turn store its KV in Valkey and a warm turn read it back.

## If the first request hangs or a replica dies

Expand All @@ -196,8 +223,9 @@ on native Linux, or `podman machine ssh 'sudo dmesg' | grep -i 'killed process'`
for a Podman machine.

Give the container runtime more memory (see [Prerequisites](#prerequisites)) or
run a single replica, then recreate the stack with `docker compose down` and
`docker compose up -d --wait`.
run the replicas one at a time (see
[Running on a smaller machine](#running-on-a-smaller-machine)), then recreate
the stack with `docker compose down` and `docker compose up -d --wait`.

## What's in this directory

Expand Down
13 changes: 8 additions & 5 deletions cookbooks/kv-caching/docker-compose.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,11 +16,12 @@
#
# Resource note: at steady state each vLLM replica uses about 3 GiB of RAM,
# the LMCache server about 0.5 GiB and Valkey about 25 MiB. Give your
# container runtime about 9 GiB for both replicas, or about 5.5 GiB to start one
# replica and skip the two-replica step (README.md, "Memory requirements", has
# the measured minimums). The batching and cache-size flags below are capped
# low because the notebook only sends one request at a time; the defaults
# over-provision for concurrency this demo never uses.
# container runtime about 9 GiB for both replicas, or about 5.5 GiB to run
# the replicas one at a time (README.md, "Running on a smaller machine"; its
# "Memory requirements" table has the measured minimums). The batching and
# cache-size flags below are capped low because the notebook only sends one
# request at a time; the defaults over-provision for concurrency this demo
# never uses.
#
# LMCache is not bundled in the vLLM CPU image, so every service that needs it
# pip-installs it on startup, pinned to one version in both places it appears
Expand Down Expand Up @@ -210,6 +211,8 @@ services:
# two-replica minimum 8.25 -> 7.5 GiB). On x86 there is no such surplus and
# starting both at once peaks no higher; the peak there is the steady state
# after both replicas have served requests.
# In the README's one-replica-at-a-time flow, start it with
# `docker compose up -d --wait --no-deps vllm-b` after stopping vllm-a.
vllm-b:
<<: *vllm-common
container_name: vllm-b
Expand Down
134 changes: 71 additions & 63 deletions cookbooks/kv-caching/kv-caching-with-valkey.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -337,69 +337,6 @@
"print(\"that prefix's KV back from Valkey instead of vLLM recomputing it.\")\n"
]
},
{
"cell_type": "markdown",
"id": "3c10668e",
"metadata": {},
"source": [
"### A second replica hits the first replica's cache\n",
"\n",
"Replica B has served no traffic and computed no KV of its own. We send it the\n",
"same conversation prefix replica A built up. The KV for that prefix lives in\n",
"Valkey, not inside replica A's process, so replica B finds it there and loads\n",
"it instead of prefilling from scratch.\n",
"\n",
"The signal: replaying replica A's prefix on replica B **reads that KV back from\n",
"Valkey** (GET calls increment) and stores **no new keys** (`dbsize` is\n",
"unchanged), because the prefix's KV is already in the shared store. Replica B\n",
"reuses what replica A computed. An in-process cache on replica A would be\n",
"invisible to replica B; putting the cache in Valkey is what makes this\n",
"cross-replica reuse possible."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "89aaa842",
"metadata": {},
"outputs": [],
"source": [
"# This is the only cell to skip when running a single replica (see \"Running\n",
"# on a smaller machine\" below). It confirms replica B is up, then replays\n",
"# replica A's exact conversation prefix (everything up to, but not\n",
"# including, replica A's last answer) as a fresh request to replica B.\n",
"print(\"replica B :\", [m.id for m in replica_b.models.list().data])\n",
"shared_prefix = messages[:-1]\n",
"\n",
"gets_before = valkey_get_calls()\n",
"before = dbsize()\n",
"resp = replica_b.chat.completions.create(\n",
" model=MODEL, messages=shared_prefix, max_tokens=32, temperature=0.0,\n",
")\n",
"after = dbsize()\n",
"gets_after = valkey_get_calls()\n",
"\n",
"print(\"Replica B answer :\", resp.choices[0].message.content)\n",
"print(\"Valkey dbsize :\", before, \"->\", after, \"(unchanged == not re-stored)\")\n",
"print(\"Valkey GET calls :\", gets_before, \"->\", gets_after, \"(went up == read from Valkey)\")\n",
"print()\n",
"print(\"Replica B served no traffic of its own, yet answered by reading the KV\")\n",
"print(\"replica A wrote to Valkey — the cache is genuinely shared through Valkey.\")\n"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"### Running on a smaller machine\n",
"\n",
"Short on memory? See\n",
"[Running on a smaller machine](README.md#running-on-a-smaller-machine-one-replica)\n",
"in the README for how to start just one replica. In that mode, skip the\n",
"cross-replica cell above (the one that calls `replica_b`) — everything else in\n",
"this notebook runs the same.\n"
]
},
{
"cell_type": "markdown",
"id": "9be5c63a",
Expand Down Expand Up @@ -462,6 +399,77 @@
" print(\"(no speedup this run — CPU timings are noisy; try re-running, and see the note above)\")\n"
]
},
{
"cell_type": "markdown",
"id": "f2a51e0c",
"metadata": {},
"source": [
"### Running on a smaller machine\n",
"\n",
"If your container runtime cannot hold both replicas at once, the README shows\n",
"two lighter ways to run this notebook (see\n",
"[Running on a smaller machine](README.md#running-on-a-smaller-machine)):\n",
"\n",
"- **One replica at a time** keeps every step, including the next one. Start\n",
" only replica A; right before the next code cell, stop replica A and start\n",
" replica B with the two `docker compose` commands from the README. Replica A's\n",
" KV stays in Valkey, so replica B still reads it back, which also shows that\n",
" the cache outlives the replica that computed it.\n",
"- **A single replica** skips the next code cell (the only one that calls\n",
" `replica_b`); everything else runs the same."
]
},
{
"cell_type": "markdown",
"id": "3c10668e",
"metadata": {},
"source": [
"### A second replica hits the first replica's cache\n",
"\n",
"Replica B has served no traffic and computed no KV of its own. We send it the\n",
"same conversation prefix replica A built up. The KV for that prefix lives in\n",
"Valkey, not inside replica A's process, so replica B finds it there and loads\n",
"it instead of prefilling from scratch.\n",
"\n",
"The signal: replaying replica A's prefix on replica B **reads that KV back from\n",
"Valkey** (GET calls increment) and stores **no new keys** (`dbsize` is\n",
"unchanged), because the prefix's KV is already in the shared store. Replica B\n",
"reuses what replica A computed. An in-process cache on replica A would be\n",
"invisible to replica B; putting the cache in Valkey is what makes this\n",
"cross-replica reuse possible."
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "89aaa842",
"metadata": {},
"outputs": [],
"source": [
"# This is the only cell to skip when running a single replica. In the\n",
"# one-replica-at-a-time flow, swap the replicas first (see \"Running on a\n",
"# smaller machine\" above). It confirms replica B is up, then replays\n",
"# replica A's exact conversation prefix (everything up to, but not\n",
"# including, replica A's last answer) as a fresh request to replica B.\n",
"print(\"replica B :\", [m.id for m in replica_b.models.list().data])\n",
"shared_prefix = messages[:-1]\n",
"\n",
"gets_before = valkey_get_calls()\n",
"before = dbsize()\n",
"resp = replica_b.chat.completions.create(\n",
" model=MODEL, messages=shared_prefix, max_tokens=32, temperature=0.0,\n",
")\n",
"after = dbsize()\n",
"gets_after = valkey_get_calls()\n",
"\n",
"print(\"Replica B answer :\", resp.choices[0].message.content)\n",
"print(\"Valkey dbsize :\", before, \"->\", after, \"(unchanged == not re-stored)\")\n",
"print(\"Valkey GET calls :\", gets_before, \"->\", gets_after, \"(went up == read from Valkey)\")\n",
"print()\n",
"print(\"Replica B served no traffic of its own, yet answered by reading the KV\")\n",
"print(\"replica A wrote to Valkey — the cache is genuinely shared through Valkey.\")\n"
]
},
{
"cell_type": "markdown",
"id": "35314202",
Expand Down