| 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 | ) |