[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>
This commit is contained in:
Jiangyun Zhu
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
f0b4b36890
commit
d648847fe4
@@ -127,7 +127,7 @@ steps:
|
||||
- tests/kernels/test_bf16_skinny_gemm.py
|
||||
commands:
|
||||
# The native NVIDIA Kimi K3 kernels require the SM100 family.
|
||||
- pytest -v -s models/kimi_k3 kernels/attention/test_kimi_k3_mla_fused_epilogue.py kernels/test_bf16_skinny_gemm.py
|
||||
- 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
|
||||
mirror:
|
||||
amd:
|
||||
label: ":amd: (MI355) Kimi K3"
|
||||
@@ -146,6 +146,26 @@ steps:
|
||||
# cuda.bindings.driver), so they are deliberately left out here.
|
||||
- pytest -v -s models/kimi_k3
|
||||
|
||||
- label: ":nvidia: (4xB200) Kimi K3 Prefix Cache Shard %N"
|
||||
key: kimi-k3-prefix-cache-4xb200
|
||||
timeout_in_minutes: 75
|
||||
device: b200-k8s
|
||||
num_devices: 4
|
||||
source_file_dependencies:
|
||||
- vllm/models/kimi_k3/
|
||||
- vllm/v1/core/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/nixl/
|
||||
- vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py
|
||||
- vllm/v1/simple_kv_offload/
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/model_executor/layers/mamba/
|
||||
- tests/models/kimi_k3/test_prefix_cache.py
|
||||
- .buildkite/scripts/install-kv-connectors.sh
|
||||
commands:
|
||||
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
|
||||
- pytest -v -s models/kimi_k3/test_prefix_cache.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT
|
||||
parallelism: 2
|
||||
|
||||
- label: ":nvidia: (H200) GLM5Next"
|
||||
key: glm5next-unit-tests
|
||||
timeout_in_minutes: 20
|
||||
|
||||
@@ -0,0 +1,375 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Kimi-K3 multi-turn prefix cache reuse with KV offload, P/D and DCP."""
|
||||
|
||||
import contextlib
|
||||
import json
|
||||
import random
|
||||
from collections.abc import Iterator
|
||||
from dataclasses import dataclass
|
||||
from math import lcm
|
||||
|
||||
import pytest
|
||||
import requests
|
||||
from prometheus_client.parser import text_string_to_metric_families
|
||||
|
||||
from tests.models.utils import check_logprobs_close
|
||||
from tests.utils import RemoteOpenAIServer, multi_gpu_marks
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
# Official Kimi-K3 config with 8 layers and no weights.
|
||||
MODEL = "riverclouds/Kimi-K3-8L-dummy"
|
||||
# The recipe's DSpark draft, reading target layers that exist in MODEL.
|
||||
DRAFT = "riverclouds/Kimi-K3-DSpark-for-8L-dummy"
|
||||
|
||||
# Conversation: the first prompt spans FIRST_PROMPT_BLOCKS blocks plus an
|
||||
# unaligned tail, and each later turn appends TURN_TOKENS random tokens.
|
||||
NUM_TURNS = 3
|
||||
FIRST_PROMPT_BLOCKS = 3
|
||||
UNALIGNED_TAIL = 37
|
||||
TURN_TOKENS = 777
|
||||
MAX_TOKENS = 8
|
||||
NUM_LOGPROBS = 5
|
||||
STARTUP_TIMEOUT = 1200
|
||||
|
||||
BASE_ARGS = [
|
||||
"--trust-remote-code",
|
||||
"--language-model-only",
|
||||
"--load-format",
|
||||
"dummy",
|
||||
"--no-enable-flashinfer-autotune",
|
||||
"--enable-prompt-tokens-details",
|
||||
# Default max_num_seqs OOMs one GPU during warmup.
|
||||
"--max-num-seqs",
|
||||
"8",
|
||||
]
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not current_platform.is_device_capability_family(100),
|
||||
reason="Kimi-K3 NVIDIA kernels require the SM100 family",
|
||||
)
|
||||
|
||||
SPEC_CONFIG = {
|
||||
"method": "dspark",
|
||||
"model": DRAFT,
|
||||
"num_speculative_tokens": 7,
|
||||
"draft_sample_method": "probabilistic",
|
||||
"rejection_sample_method": "block",
|
||||
}
|
||||
|
||||
NIXL = {"kv_connector": "NixlConnector", "kv_role": "kv_both"}
|
||||
# SimpleCPUOffloadConnector as in the Kimi-K3 recipe.
|
||||
OFFLOAD = {
|
||||
"kv_connector": "SimpleCPUOffloadConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_connector_extra_config": {
|
||||
"cpu_bytes_to_use_per_rank": 32 << 30,
|
||||
"lazy_offload": False,
|
||||
},
|
||||
}
|
||||
NIXL_OFFLOAD = {
|
||||
"kv_connector": "MultiConnector",
|
||||
"kv_role": "kv_both",
|
||||
"kv_connector_extra_config": {"connectors": [NIXL, OFFLOAD]},
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Mode:
|
||||
# A finer prefix_match_unit enables partial hits inside a Mamba block.
|
||||
prefix_match_unit: int | None
|
||||
spec: bool
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
name = "block" if self.prefix_match_unit is None else "partial"
|
||||
return f"{name}-dspark" if self.spec else name
|
||||
|
||||
|
||||
MODES = [Mode(unit, spec) for spec in (False, True) for unit in (None, 128)]
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Instance:
|
||||
tp: int = 1
|
||||
dp: int = 1
|
||||
dcp: int = 1
|
||||
kv_config: dict | None = None
|
||||
prefix_caching: bool = True
|
||||
|
||||
@property
|
||||
def num_gpus(self) -> int:
|
||||
return self.tp * self.dp
|
||||
|
||||
def args(self, mode: Mode) -> list[str]:
|
||||
args = BASE_ARGS + ["--tensor-parallel-size", str(self.tp)]
|
||||
if mode.spec:
|
||||
args += ["--speculative-config", json.dumps(SPEC_CONFIG)]
|
||||
if self.dp > 1:
|
||||
args += ["--data-parallel-size", str(self.dp), "--enable-expert-parallel"]
|
||||
if self.dcp > 1:
|
||||
args += ["--decode-context-parallel-size", str(self.dcp)]
|
||||
if self.kv_config is not None:
|
||||
args += ["--kv-transfer-config", json.dumps(self.kv_config)]
|
||||
if not self.prefix_caching:
|
||||
return args + ["--no-enable-prefix-caching"]
|
||||
args.append("--enable-prefix-caching")
|
||||
if mode.prefix_match_unit is not None:
|
||||
args += ["--prefix-match-unit", str(mode.prefix_match_unit)]
|
||||
return args
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Deployment:
|
||||
decode: Instance
|
||||
prefill: Instance | None = None
|
||||
# Reset GPU prefix caches each turn so hits must come from CPU.
|
||||
offload: bool = False
|
||||
|
||||
@property
|
||||
def instances(self) -> tuple[Instance, ...]:
|
||||
return (self.prefill, self.decode) if self.prefill else (self.decode,)
|
||||
|
||||
@property
|
||||
def num_gpus(self) -> int:
|
||||
return sum(i.num_gpus for i in self.instances)
|
||||
|
||||
|
||||
DEPLOYMENTS = {
|
||||
"plain": Deployment(Instance()),
|
||||
"offload": Deployment(Instance(kv_config=OFFLOAD), offload=True),
|
||||
"dcp2": Deployment(Instance(tp=2, dcp=2)),
|
||||
"dcp2-offload": Deployment(Instance(tp=2, dcp=2, kv_config=OFFLOAD), offload=True),
|
||||
"pd": Deployment(Instance(kv_config=NIXL), prefill=Instance(kv_config=NIXL)),
|
||||
# NIXL rejects prefix caching on a hybrid decoder with a different TP.
|
||||
"pd-tp2-dep2": Deployment(
|
||||
Instance(dp=2, kv_config=NIXL, prefix_caching=False),
|
||||
prefill=Instance(tp=2, kv_config=NIXL),
|
||||
),
|
||||
"pd-offload": Deployment(
|
||||
Instance(kv_config=NIXL_OFFLOAD),
|
||||
prefill=Instance(kv_config=NIXL_OFFLOAD),
|
||||
offload=True,
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def _known_failure(name: str, mode: Mode) -> str | None:
|
||||
deployment = DEPLOYMENTS[name]
|
||||
if mode.spec and deployment.decode.dcp > 1:
|
||||
return "FlashInfer MLA DCP decode breaks on ragged spec batches (#59392)"
|
||||
return None
|
||||
|
||||
|
||||
def _case(name: str, mode: Mode):
|
||||
marks = []
|
||||
if (num_gpus := DEPLOYMENTS[name].num_gpus) > 1:
|
||||
marks += multi_gpu_marks(num_gpus=num_gpus)
|
||||
if reason := _known_failure(name, mode):
|
||||
marks.append(pytest.mark.xfail(reason=reason, strict=True))
|
||||
return pytest.param(name, mode, marks=marks, id=f"{name}-{mode.name}")
|
||||
|
||||
|
||||
def _block_size(url: str) -> int:
|
||||
response = requests.get(f"{url}/metrics", timeout=30)
|
||||
response.raise_for_status()
|
||||
for family in text_string_to_metric_families(response.text):
|
||||
for sample in family.samples:
|
||||
if sample.name == "vllm:cache_config_info":
|
||||
return lcm(
|
||||
*(
|
||||
int(value)
|
||||
for key in ("block_size", "mamba_block_size")
|
||||
if (value := sample.labels.get(key, "None")) != "None"
|
||||
)
|
||||
)
|
||||
raise AssertionError("missing vllm:cache_config_info metric")
|
||||
|
||||
|
||||
def _complete(url: str, prompt: list[int], salt: str, **extra) -> dict:
|
||||
body = {
|
||||
"model": MODEL,
|
||||
"prompt": prompt,
|
||||
"max_tokens": MAX_TOKENS,
|
||||
"temperature": 0,
|
||||
"logprobs": NUM_LOGPROBS,
|
||||
"return_tokens_as_token_ids": True,
|
||||
"cache_salt": salt,
|
||||
**extra,
|
||||
}
|
||||
response = requests.post(f"{url}/v1/completions", json=body, timeout=600)
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
|
||||
|
||||
def _pd_complete(
|
||||
prefill_url: str, decode_url: str, prompt: list[int], salt: str
|
||||
) -> tuple[dict, dict]:
|
||||
remote_decode = {
|
||||
"do_remote_decode": True,
|
||||
"do_remote_prefill": False,
|
||||
"remote_engine_id": None,
|
||||
"remote_block_ids": None,
|
||||
"remote_host": None,
|
||||
"remote_port": None,
|
||||
}
|
||||
prefill = _complete(
|
||||
prefill_url, prompt, salt, max_tokens=1, kv_transfer_params=remote_decode
|
||||
)
|
||||
decode = _complete(
|
||||
decode_url, prompt, salt, kv_transfer_params=prefill["kv_transfer_params"]
|
||||
)
|
||||
return prefill, decode
|
||||
|
||||
|
||||
def _cached(response: dict) -> int:
|
||||
return response["usage"]["prompt_tokens_details"]["cached_tokens"]
|
||||
|
||||
|
||||
def _token_id(token: str) -> int:
|
||||
return int(token.removeprefix("token_id:"))
|
||||
|
||||
|
||||
def _tokens_text_logprobs(response: dict):
|
||||
choice = response["choices"][0]
|
||||
ids = [_token_id(t) for t in choice["logprobs"]["tokens"]]
|
||||
top = [
|
||||
{_token_id(t): lp for t, lp in d.items()}
|
||||
for d in choice["logprobs"]["top_logprobs"]
|
||||
]
|
||||
return ids, choice["text"], top
|
||||
|
||||
|
||||
def _reset_gpu_prefix_cache(url: str) -> None:
|
||||
response = requests.post(
|
||||
f"{url}/reset_prefix_cache", params={"reset_external": "false"}, timeout=60
|
||||
)
|
||||
response.raise_for_status()
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Turn:
|
||||
prompt_len: int
|
||||
cached: int
|
||||
expected_cached: int
|
||||
recompute_cached: int
|
||||
output: tuple
|
||||
recompute: tuple
|
||||
|
||||
@property
|
||||
def matches_recompute(self) -> bool:
|
||||
return self.output[0] == self.recompute[0]
|
||||
|
||||
|
||||
def _expected_cached(prev_prompt_len: int, hit_unit: int, spec: bool) -> int:
|
||||
"""A turn must reuse the previous prompt up to its last checkpoint.
|
||||
|
||||
EAGLE-style drafts drop the trailing unit, so the checkpoint is one earlier.
|
||||
"""
|
||||
if not prev_prompt_len:
|
||||
return 0
|
||||
checkpoint = (prev_prompt_len - 1) // hit_unit * hit_unit
|
||||
return max(checkpoint - hit_unit, 0) if spec else checkpoint
|
||||
|
||||
|
||||
@contextlib.contextmanager
|
||||
def _serve(deployment: Deployment, mode: Mode) -> Iterator[list[RemoteOpenAIServer]]:
|
||||
with contextlib.ExitStack() as stack:
|
||||
# Sequential shutdown waits on memory still held by sibling servers.
|
||||
servers: list[RemoteOpenAIServer] = []
|
||||
stack.callback(RemoteOpenAIServer.shutdown_many, servers)
|
||||
next_gpu = 0
|
||||
for instance in deployment.instances:
|
||||
env = {}
|
||||
if deployment.offload:
|
||||
# Enables /reset_prefix_cache.
|
||||
env["VLLM_SERVER_DEV_MODE"] = "1"
|
||||
if deployment.prefill is not None:
|
||||
gpus = range(next_gpu, next_gpu + instance.num_gpus)
|
||||
next_gpu += instance.num_gpus
|
||||
env["CUDA_VISIBLE_DEVICES"] = ",".join(map(str, gpus))
|
||||
env["VLLM_SSM_CONV_STATE_LAYOUT"] = "DS"
|
||||
env["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(get_open_port())
|
||||
servers.append(
|
||||
RemoteOpenAIServer(
|
||||
MODEL,
|
||||
instance.args(mode),
|
||||
env_dict=env,
|
||||
max_wait_seconds=STARTUP_TIMEOUT,
|
||||
)
|
||||
)
|
||||
yield servers
|
||||
|
||||
|
||||
def _run_conversation(
|
||||
deployment: Deployment,
|
||||
servers: list[RemoteOpenAIServer],
|
||||
mode: Mode,
|
||||
) -> list[Turn]:
|
||||
# Under P/D the prefiller computes the prompt, so hits are read from it.
|
||||
compute_url, decode_url = servers[0].url_root, servers[-1].url_root
|
||||
block_size = _block_size(compute_url)
|
||||
hit_unit = mode.prefix_match_unit or block_size
|
||||
rng = random.Random(0)
|
||||
|
||||
def new_tokens(n: int) -> list[int]:
|
||||
return [rng.randint(1000, 150000) for _ in range(n)]
|
||||
|
||||
prompt = new_tokens(FIRST_PROMPT_BLOCKS * block_size + UNALIGNED_TAIL)
|
||||
prev_len = 0
|
||||
turns = []
|
||||
for turn in range(NUM_TURNS):
|
||||
if deployment.offload:
|
||||
for server in servers:
|
||||
_reset_gpu_prefix_cache(server.url_root)
|
||||
if deployment.prefill is not None:
|
||||
computed, output = _pd_complete(compute_url, decode_url, prompt, "conv")
|
||||
else:
|
||||
computed = output = _complete(decode_url, prompt, "conv")
|
||||
recompute = _complete(decode_url, prompt, f"recompute-{turn}")
|
||||
turns.append(
|
||||
Turn(
|
||||
prompt_len=len(prompt),
|
||||
cached=_cached(computed),
|
||||
expected_cached=_expected_cached(prev_len, hit_unit, mode.spec),
|
||||
recompute_cached=_cached(recompute),
|
||||
output=_tokens_text_logprobs(output),
|
||||
recompute=_tokens_text_logprobs(recompute),
|
||||
)
|
||||
)
|
||||
prev_len = len(prompt)
|
||||
prompt = prompt + new_tokens(TURN_TOKENS)
|
||||
return turns
|
||||
|
||||
|
||||
def _format(turns: list[Turn]) -> str:
|
||||
rows = ["turn prompt cached expected matches_recompute"]
|
||||
rows += [
|
||||
f"{i:>4} {t.prompt_len:>6} {t.cached:>6} {t.expected_cached:>8} "
|
||||
f"{t.matches_recompute}"
|
||||
for i, t in enumerate(turns)
|
||||
]
|
||||
return "\n".join(rows)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"name, mode", [_case(name, mode) for name in DEPLOYMENTS for mode in MODES]
|
||||
)
|
||||
def test_turns_reuse_prefix_and_match_recompute(name: str, mode: Mode) -> None:
|
||||
deployment = DEPLOYMENTS[name]
|
||||
with _serve(deployment, mode) as servers:
|
||||
turns = _run_conversation(deployment, servers, mode)
|
||||
table = _format(turns)
|
||||
print(f"\n{name}-{mode.name}\n{table}")
|
||||
|
||||
assert all(t.recompute_cached == 0 for t in turns), table
|
||||
assert all(t.cached >= t.expected_cached for t in turns), table
|
||||
# Dummy weights differ across TP sizes.
|
||||
if len({i.tp for i in deployment.instances}) == 1:
|
||||
check_logprobs_close(
|
||||
outputs_0_lst=[t.recompute for t in turns],
|
||||
outputs_1_lst=[t.output for t in turns],
|
||||
name_0="recompute",
|
||||
name_1=name,
|
||||
)
|
||||
Reference in New Issue
Block a user