[CI][Kimi-K3] Test prefix cache reuse with KV offload, P/D and DCP (#59229)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Jiangyun Zhu committed
Oct 2, 2026 at 20:58 UTC
d648847fe451a6a0366f8dd0e1ae96804a49a299
2 files changed
+396
-1
.buildkite/test_areas/models_basic.yaml
+21
-1
@@ -127,7 +127,7 @@ steps:
127
- tests/kernels/test_bf16_skinny_gemm.py
128
commands:
129
# The native NVIDIA Kimi K3 kernels require the SM100 family.
130
- - pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
130
+ - pytest -v -s models/kimi_k3 --ignore=models/kimi_k3/test_prefix_cache.py kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
131
mirror:
132
amd:
133
label: ":amd: (MI355) Kimi K3"
@@ -146,6 +146,26 @@ steps:
146
# cuda.bindings.driver), so they are deliberately left out here.
147
- pytest -v -s models/kimi_k3
148
149
+- label: ":nvidia: (4xB200) Kimi K3 Prefix Cache Shard %N"
150
+ key: kimi-k3-prefix-cache-4xb200
151
+ timeout_in_minutes: 75
152
+ device: b200-k8s
153
+ num_devices: 4
154
+ source_file_dependencies:
155
+ - vllm/models/kimi_k3/
156
+ - vllm/v1/core/
157
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
158
+ - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
159
+ - vllm/v1/simple_kv_offload/
160
+ - vllm/v1/spec_decode/
161
+ - vllm/model_executor/layers/mamba/
162
+ - tests/models/kimi_k3/test_prefix_cache.py
163
+ - .buildkite/scripts/install-kv-connectors.sh
164
+ commands:
165
+ - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
166
+ - pytest -v -s models/kimi_k3/test_prefix_cache.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
167
+ parallelism: 2
168
+
169
- label: ":nvidia: (H200) GLM5Next"
170
key: glm5next-unit-tests
171
timeout_in_minutes: 20
tests/models/kimi_k3/test_prefix_cache.py
new
+375
@@ -0,0 +1,375 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Kimi-K3 multi-turn prefix cache reuse with KV offload, P/D and DCP."""
4
+
5
+import contextlib
6
+import json
7
+import random
8
+from collections.abc import Iterator
9
+from dataclasses import dataclass
10
+from math import lcm
11
+
12
+import pytest
13
+import requests
14
+from prometheus_client.parser import text_string_to_metric_families
15
+
16
+from tests.models.utils import check_logprobs_close
17
+from tests.utils import RemoteOpenAIServer, multi_gpu_marks
18
+from vllm.platforms import current_platform
19
+from vllm.utils.network_utils import get_open_port
20
+
21
+# Official Kimi-K3 config with 8 layers and no weights.
22
+MODEL = "riverclouds/Kimi-K3-8L-dummy"
23
+# The recipe's DSpark draft, reading target layers that exist in MODEL.
24
+DRAFT = "riverclouds/Kimi-K3-DSpark-for-8L-dummy"
25
+
26
+# Conversation: the first prompt spans FIRST_PROMPT_BLOCKS blocks plus an
27
+# unaligned tail, and each later turn appends TURN_TOKENS random tokens.
28
+NUM_TURNS = 3
29
+FIRST_PROMPT_BLOCKS = 3
30
+UNALIGNED_TAIL = 37
31
+TURN_TOKENS = 777
32
+MAX_TOKENS = 8
33
+NUM_LOGPROBS = 5
34
+STARTUP_TIMEOUT = 1200
35
+
36
+BASE_ARGS = [
37
+ "--trust-remote-code",
38
+ "--language-model-only",
39
+ "--load-format",
40
+ "dummy",
41
+ "--no-enable-flashinfer-autotune",
42
+ "--enable-prompt-tokens-details",
43
+ # Default max_num_seqs OOMs one GPU during warmup.
44
+ "--max-num-seqs",
45
+ "8",
46
+]
47
+
48
+pytestmark = pytest.mark.skipif(
49
+ not current_platform.is_device_capability_family(100),
50
+ reason="Kimi-K3 NVIDIA kernels require the SM100 family",
51
+)
52
+
53
+SPEC_CONFIG = {
54
+ "method": "dspark",
55
+ "model": DRAFT,
56
+ "num_speculative_tokens": 7,
57
+ "draft_sample_method": "probabilistic",
58
+ "rejection_sample_method": "block",
59
+}
60
+
61
+NIXL = {"kv_connector": "NixlConnector", "kv_role": "kv_both"}
62
+# SimpleCPUOffloadConnector as in the Kimi-K3 recipe.
63
+OFFLOAD = {
64
+ "kv_connector": "SimpleCPUOffloadConnector",
65
+ "kv_role": "kv_both",
66
+ "kv_connector_extra_config": {
67
+ "cpu_bytes_to_use_per_rank": 32 << 30,
68
+ "lazy_offload": False,
69
+ },
70
+}
71
+NIXL_OFFLOAD = {
72
+ "kv_connector": "MultiConnector",
73
+ "kv_role": "kv_both",
74
+ "kv_connector_extra_config": {"connectors": [NIXL, OFFLOAD]},
75
+}
76
+
77
+
78
+@dataclass(frozen=True)
79
+class Mode:
80
+ # A finer prefix_match_unit enables partial hits inside a Mamba block.
81
+ prefix_match_unit: int | None
82
+ spec: bool
83
+
84
+ @property
85
+ def name(self) -> str:
86
+ name = "block" if self.prefix_match_unit is None else "partial"
87
+ return f"{name}-dspark" if self.spec else name
88
+
89
+
90
+MODES = [Mode(unit, spec) for spec in (False, True) for unit in (None, 128)]
91
+
92
+
93
+@dataclass(frozen=True)
94
+class Instance:
95
+ tp: int = 1
96
+ dp: int = 1
97
+ dcp: int = 1
98
+ kv_config: dict | None = None
99
+ prefix_caching: bool = True
100
+
101
+ @property
102
+ def num_gpus(self) -> int:
103
+ return self.tp * self.dp
104
+
105
+ def args(self, mode: Mode) -> list[str]:
106
+ args = BASE_ARGS + ["--tensor-parallel-size", str(self.tp)]
107
+ if mode.spec:
108
+ args += ["--speculative-config", json.dumps(SPEC_CONFIG)]
109
+ if self.dp > 1:
110
+ args += ["--data-parallel-size", str(self.dp), "--enable-expert-parallel"]
111
+ if self.dcp > 1:
112
+ args += ["--decode-context-parallel-size", str(self.dcp)]
113
+ if self.kv_config is not None:
114
+ args += ["--kv-transfer-config", json.dumps(self.kv_config)]
115
+ if not self.prefix_caching:
116
+ return args + ["--no-enable-prefix-caching"]
117
+ args.append("--enable-prefix-caching")
118
+ if mode.prefix_match_unit is not None:
119
+ args += ["--prefix-match-unit", str(mode.prefix_match_unit)]
120
+ return args
121
+
122
+
123
+@dataclass(frozen=True)
124
+class Deployment:
125
+ decode: Instance
126
+ prefill: Instance | None = None
127
+ # Reset GPU prefix caches each turn so hits must come from CPU.
128
+ offload: bool = False
129
+
130
+ @property
131
+ def instances(self) -> tuple[Instance, ...]:
132
+ return (self.prefill, self.decode) if self.prefill else (self.decode,)
133
+
134
+ @property
135
+ def num_gpus(self) -> int:
136
+ return sum(i.num_gpus for i in self.instances)
137
+
138
+
139
+DEPLOYMENTS = {
140
+ "plain": Deployment(Instance()),
141
+ "offload": Deployment(Instance(kv_config=OFFLOAD), offload=True),
142
+ "dcp2": Deployment(Instance(tp=2, dcp=2)),
143
+ "dcp2-offload": Deployment(Instance(tp=2, dcp=2, kv_config=OFFLOAD), offload=True),
144
+ "pd": Deployment(Instance(kv_config=NIXL), prefill=Instance(kv_config=NIXL)),
145
+ # NIXL rejects prefix caching on a hybrid decoder with a different TP.
146
+ "pd-tp2-dep2": Deployment(
147
+ Instance(dp=2, kv_config=NIXL, prefix_caching=False),
148
+ prefill=Instance(tp=2, kv_config=NIXL),
149
+ ),
150
+ "pd-offload": Deployment(
151
+ Instance(kv_config=NIXL_OFFLOAD),
152
+ prefill=Instance(kv_config=NIXL_OFFLOAD),
153
+ offload=True,
154
+ ),
155
+}
156
+
157
+
158
+def _known_failure(name: str, mode: Mode) -> str | None:
159
+ deployment = DEPLOYMENTS[name]
160
+ if mode.spec and deployment.decode.dcp > 1:
161
+ return "FlashInfer MLA DCP decode breaks on ragged spec batches (#59392)"
162
+ return None
163
+
164
+
165
+def _case(name: str, mode: Mode):
166
+ marks = []
167
+ if (num_gpus := DEPLOYMENTS[name].num_gpus) > 1:
168
+ marks += multi_gpu_marks(num_gpus=num_gpus)
169
+ if reason := _known_failure(name, mode):
170
+ marks.append(pytest.mark.xfail(reason=reason, strict=True))
171
+ return pytest.param(name, mode, marks=marks, id=f"{name}-{mode.name}")
172
+
173
+
174
+def _block_size(url: str) -> int:
175
+ response = requests.get(f"{url}/metrics", timeout=30)
176
+ response.raise_for_status()
177
+ for family in text_string_to_metric_families(response.text):
178
+ for sample in family.samples:
179
+ if sample.name == "vllm:cache_config_info":
180
+ return lcm(
181
+ *(
182
+ int(value)
183
+ for key in ("block_size", "mamba_block_size")
184
+ if (value := sample.labels.get(key, "None")) != "None"
185
+ )
186
+ )
187
+ raise AssertionError("missing vllm:cache_config_info metric")
188
+
189
+
190
+def _complete(url: str, prompt: list[int], salt: str, **extra) -> dict:
191
+ body = {
192
+ "model": MODEL,
193
+ "prompt": prompt,
194
+ "max_tokens": MAX_TOKENS,
195
+ "temperature": 0,
196
+ "logprobs": NUM_LOGPROBS,
197
+ "return_tokens_as_token_ids": True,
198
+ "cache_salt": salt,
199
+ **extra,
200
+ }
201
+ response = requests.post(f"{url}/v1/completions", json=body, timeout=600)
202
+ response.raise_for_status()
203
+ return response.json()
204
+
205
+
206
+def _pd_complete(
207
+ prefill_url: str, decode_url: str, prompt: list[int], salt: str
208
+) -> tuple[dict, dict]:
209
+ remote_decode = {
210
+ "do_remote_decode": True,
211
+ "do_remote_prefill": False,
212
+ "remote_engine_id": None,
213
+ "remote_block_ids": None,
214
+ "remote_host": None,
215
+ "remote_port": None,
216
+ }
217
+ prefill = _complete(
218
+ prefill_url, prompt, salt, max_tokens=1, kv_transfer_params=remote_decode
219
+ )
220
+ decode = _complete(
221
+ decode_url, prompt, salt, kv_transfer_params=prefill["kv_transfer_params"]
222
+ )
223
+ return prefill, decode
224
+
225
+
226
+def _cached(response: dict) -> int:
227
+ return response["usage"]["prompt_tokens_details"]["cached_tokens"]
228
+
229
+
230
+def _token_id(token: str) -> int:
231
+ return int(token.removeprefix("token_id:"))
232
+
233
+
234
+def _tokens_text_logprobs(response: dict):
235
+ choice = response["choices"][0]
236
+ ids = [_token_id(t) for t in choice["logprobs"]["tokens"]]
237
+ top = [
238
+ {_token_id(t): lp for t, lp in d.items()}
239
+ for d in choice["logprobs"]["top_logprobs"]
240
+ ]
241
+ return ids, choice["text"], top
242
+
243
+
244
+def _reset_gpu_prefix_cache(url: str) -> None:
245
+ response = requests.post(
246
+ f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60
247
+ )
248
+ response.raise_for_status()
249
+
250
+
251
+@dataclass(frozen=True)
252
+class Turn:
253
+ prompt_len: int
254
+ cached: int
255
+ expected_cached: int
256
+ recompute_cached: int
257
+ output: tuple
258
+ recompute: tuple
259
+
260
+ @property
261
+ def matches_recompute(self) -> bool:
262
+ return self.output[0] == self.recompute[0]
263
+
264
+
265
+def _expected_cached(prev_prompt_len: int, hit_unit: int, spec: bool) -> int:
266
+ """A turn must reuse the previous prompt up to its last checkpoint.
267
+
268
+ EAGLE-style drafts drop the trailing unit, so the checkpoint is one earlier.
269
+ """
270
+ if not prev_prompt_len:
271
+ return 0
272
+ checkpoint = (prev_prompt_len - 1) // hit_unit * hit_unit
273
+ return max(checkpoint - hit_unit, 0) if spec else checkpoint
274
+
275
+
276
+@contextlib.contextmanager
277
+def _serve(deployment: Deployment, mode: Mode) -> Iterator[list[RemoteOpenAIServer]]:
278
+ with contextlib.ExitStack() as stack:
279
+ # Sequential shutdown waits on memory still held by sibling servers.
280
+ servers: list[RemoteOpenAIServer] = []
281
+ stack.callback(RemoteOpenAIServer.shutdown_many, servers)
282
+ next_gpu = 0
283
+ for instance in deployment.instances:
284
+ env = {}
285
+ if deployment.offload:
286
+ # Enables /reset_prefix_cache.
287
+ env["VLLM_SERVER_DEV_MODE"] = "1"
288
+ if deployment.prefill is not None:
289
+ gpus = range(next_gpu, next_gpu + instance.num_gpus)
290
+ next_gpu += instance.num_gpus
291
+ env["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpus))
292
+ env["VLLM_SSM_CONV_STATE_LAYOUT"] = "DS"
293
+ env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(get_open_port())
294
+ servers.append(
295
+ RemoteOpenAIServer(
296
+ MODEL,
297
+ instance.args(mode),
298
+ env_dict=env,
299
+ max_wait_seconds=STARTUP_TIMEOUT,
300
+ )
301
+ )
302
+ yield servers
303
+
304
+
305
+def _run_conversation(
306
+ deployment: Deployment,
307
+ servers: list[RemoteOpenAIServer],
308
+ mode: Mode,
309
+) -> list[Turn]:
310
+ # Under P/D the prefiller computes the prompt, so hits are read from it.
311
+ compute_url, decode_url = servers[0].url_root, servers[-1].url_root
312
+ block_size = _block_size(compute_url)
313
+ hit_unit = mode.prefix_match_unit or block_size
314
+ rng = random.Random(0)
315
+
316
+ def new_tokens(n: int) -> list[int]:
317
+ return [rng.randint(1000, 150000) for _ in range(n)]
318
+
319
+ prompt = new_tokens(FIRST_PROMPT_BLOCKS * block_size + UNALIGNED_TAIL)
320
+ prev_len = 0
321
+ turns = []
322
+ for turn in range(NUM_TURNS):
323
+ if deployment.offload:
324
+ for server in servers:
325
+ _reset_gpu_prefix_cache(server.url_root)
326
+ if deployment.prefill is not None:
327
+ computed, output = _pd_complete(compute_url, decode_url, prompt, "conv")
328
+ else:
329
+ computed = output = _complete(decode_url, prompt, "conv")
330
+ recompute = _complete(decode_url, prompt, f"recompute-{turn}")
331
+ turns.append(
332
+ Turn(
333
+ prompt_len=len(prompt),
334
+ cached=_cached(computed),
335
+ expected_cached=_expected_cached(prev_len, hit_unit, mode.spec),
336
+ recompute_cached=_cached(recompute),
337
+ output=_tokens_text_logprobs(output),
338
+ recompute=_tokens_text_logprobs(recompute),
339
+ )
340
+ )
341
+ prev_len = len(prompt)
342
+ prompt = prompt + new_tokens(TURN_TOKENS)
343
+ return turns
344
+
345
+
346
+def _format(turns: list[Turn]) -> str:
347
+ rows = ["turn prompt cached expected matches_recompute"]
348
+ rows += [
349
+ f"{i:>4} {t.prompt_len:>6} {t.cached:>6} {t.expected_cached:>8} "
350
+ f"{t.matches_recompute}"
351
+ for i, t in enumerate(turns)
352
+ ]
353
+ return "\n".join(rows)
354
+
355
+
356
+@pytest.mark.parametrize(
357
+ "name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES]
358
+)
359
+def test_turns_reuse_prefix_and_match_recompute(name: str, mode: Mode) -> None:
360
+ deployment = DEPLOYMENTS[name]
361
+ with _serve(deployment, mode) as servers:
362
+ turns = _run_conversation(deployment, servers, mode)
363
+ table = _format(turns)
364
+ print(f"\n{name}-{mode.name}\n{table}")
365
+
366
+ assert all(t.recompute_cached == 0 for t in turns), table
367
+ assert all(t.cached >= t.expected_cached for t in turns), table
368
+ # Dummy weights differ across TP sizes.
369
+ if len({i.tp for i in deployment.instances}) == 1:
370
+ check_logprobs_close(
371
+ outputs_0_lst=[t.recompute for t in turns],
372
+ outputs_1_lst=[t.output for t in turns],
373
+ name_0="recompute",
374
+ name_1=name,
375
+ )