Prhub

#46844 [CI] Mooncake PD integration tests

原始 PR 作者 NickLucche 合并时间 2026-08-03 21:55 文件变更 9 提交数 2 评论 10 代码增减 +413 / -4

执行摘要

为 Mooncake PD 引入端到端集成测试与独立 CI 任务

PR body 明确说明出发点:"Let's start tracking MC functionalities e2e with proper integration tests, starting with basic PD functionality"。作者把 Buildkite 任务拆分为独立组,是因为 Mooncake 的 feature matrix compatibility 尚未达到 NixlConnector 的成熟度,需要能单独开关、独立演进并节省 CI 开销。Issue 评论记录了前置条件:等待 CUDA 13 预编译 wheel 发布("Waiting on cu13 precompiled package to ship"),以及 RDMA 依赖 CI 节点分配的问题("Having some issues on RDMA depending on node assignment"),这些直接塑造了本 PR 的池选择与 wheel 替换设计。

值得精读,尤其是 run_accuracy_test.sh 中 RDMA 设备白名单的注释、.buildkite/test_areas/disaggregated_mooncake.yaml 里池选择的理由、以及 install-kv-connectors.sh 的 CUDA 13 wheel 替换逻辑——这三处都是真实 CI 环境踩坑后的产物,对想要把自有 P/D 或 RDMA 传输测试接入 CI 的团队有直接参考价值。整体评分定位在"有意义的测试基建"而非核心功能变更。

讨论亮点

zhewenl 作为主要 reviewer 提出两点:一是质疑测试目录里新增的 toy_proxy_server.py 与示例代理重复("Why couldn't we just reuse the example proxy?"),NickLucche 承认 "I overfitted on nixl setup lol",最终版本删除了 toy 代理、直接复用 examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py;二是建议补充 SWA 模型覆盖("maybe https://huggingface.co/google/gemma-3-1b-it will be a good fit"),NickLucche 回应这是 "a PR to get things started",当前重点是打通基础设施(linking errors 等),不同模型配置可仿照 nixl 的做法后续按模型类型补充,该建议目前仍是开放项。ivanium 附议 "Overall LGTM and agree with @zhewenl's comments",zhewenl 最终 APPROVED。

实现拆解

实现分为五步:

  1. 新增测试包与精度断言tests/v1/kv_connector/mooncake_integration/__init__.py 使目录成为可导入的测试包;test_accuracy.py 定义 run_simple_prompt()(经 OpenAI 兼容接口向本地 8192 端口代理发一个续写请求做冒烟验证)与 test_accuracy()(用 lm_eval 的 local-completions 后端跑 gsm8k,以 exact_match,strict-match 指标与 EXPECTED_VALUES 中 Qwen/Qwen3-0.6B 的 0.41 基线做 ±5% 容差断言)。未登记的模型只告警不失败,避免 CI 误报。

  2. 编排脚本 run_accuracy_test.sh。全部参数通过环境变量注入:NUM_PREFILL_INSTANCESNUM_DECODE_INSTANCESPREFILLER_TP_SIZEDECODER_TP_SIZEGPU_MEMORY_UTILIZATIONBLOCK_SIZEVLLM_SERVE_EXTRA_ARGSATTENTION_BACKENDkv_config() 生成 kv_producer/kv_consumer 两份 --kv-transfer-config;端口规划为 prefill 8100+、decode 8200+、bootstrap 8998+,GPU 按实例序号轮转分配。所有实例 wait_for_server 就绪后,复用官方 mooncake_connector_proxy.py 聚合出 8192 端口入口,最后执行 pytest,并由 cleanup_instances() 清理残留 vllm serve 进程。MOONCAKE_DEVICE_NAME 默认 mlx5_12 是对 b200-k8s 池多 HCA 拓扑踩坑后的白名单方案。

  3. 配置扫描 config_sweep_accuracy_test.sh。先执行 test_mooncake_imports.py 作为 import canary,再用 env + bash 数组方式安全传参(避免 eval),依次扫 GPU_MEMORY_UTILIZATION=0.6 下 TP 1x1 与 2x2 两种对称配置;支持 ATTENTION_BACKEND 透传,任一配置失败即 exit 1

  4. 独立 Buildkite 任务组.buildkite/test_areas/disaggregated_mooncake.yaml 新增 "Disaggregated Mooncake" 组,device: b200-k8snum_devices: 4timeout_in_minutes: 10。选择该池的唯一理由是它能承载 Mooncake 传输:只有它授予 IPC_LOCK(RDMA 钉页必需)并使用 host networking;其他池经拓扑发现找不到可用 HCA,tcp fallback 会广告 docker bridge 地址导致连接全部失败。source_file_dependencies 限定仅 Mooncake 连接器源码与集成测试目录变更时触发。

  5. 依赖与安装脚本配套requirements/kv_connectors.txtmooncake-transfer-engine 提升到 >= 0.3.12 并注明 CUDA 12/13 双 wheel 关系;install-kv-connectors.sh 的元数据探测从只查 nixl 扩展为同时输出 mooncake 版本,并在 CUDA 13 镜像上先卸载 CUDA 12 wheel 再安装同版本 mooncake-transfer-engine-cuda13test_mooncake_imports.py 从包元数据层提前暴露 wheel 不匹配,避免引擎启动才出现 "Mooncake is not available"。docs/features/mooncake_connector_usage.md 补充了 device_name 白名单参数文档。

文件 模块 状态 重要度
tests/v1/kv_connector/mooncake_integration/test_accuracy.py 精度测试 added 6.57
tests/v1/kv_connector/mooncake_integration/run_accuracy_test.sh 编排脚本 added 6.53
tests/v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh 配置扫描 added 5.63
tests/v1/kv_connector/mooncake_integration/test_mooncake_imports.py 导入检查 added 5.16
.buildkite/test_areas/disaggregated_mooncake.yaml CI 任务 added 4.4
.buildkite/scripts/install-kv-connectors.sh 安装脚本 modified 4.56
tests/v1/kv_connector/mooncake_integration/__init__.py 测试包 added 2.31
requirements/kv_connectors.txt 依赖清单 modified 1.9
docs/features/mooncake_connector_usage.md 使用文档 modified 1.58

关键符号

run_simple_prompt test_accuracy test_mooncake_engine_imports run_tests_for_model run_tests kv_config wait_for_server cleanup_instances

关键源码片段

tests/v1/kv_connector/mooncake_integration/test_accuracy.py test-coverage

端到端精度断言的核心测试文件:先做冒烟请求,再用 lm_eval 跑 gsm8k 与模型基线做 ±5% 容差比对,是整条 P/D 链路正确性的最终验收点。

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import osimport lm_eval
import openaiBASE_URL = "http://localhost:8192/v1"
NUM_CONCURRENT = 100 # lm_eval 并发请求数
TASK = "gsm8k"
FILTER = "exact_match,strict-match"
RTOL = 0.05 # 相对容差:允许精度波动 5%# 模型对应的期望精度基线,来自本地 benchmark 实测。
EXPECTED_VALUES = {
    "Qwen/Qwen3-0.6B": 0.41,
}SIMPLE_PROMPT = (
    "The best part about working on vLLM is that I got to meet so many people across "
    "various different organizations like UCB, Google, and Meta which means",
)# 通过环境变量覆盖被测模型,默认 Qwen/Qwen3-0.6B。
MODEL_NAME = os.environ.get("TEST_MODEL", "Qwen/Qwen3-0.6B")
​
​
def run_simple_prompt():
    """冒烟请求:确认 P/D 两侧经 Mooncake 转移 KV 后仍能正常续写。"""
    client = openai.OpenAI(api_key="EMPTY", base_url=BASE_URL)
    completion = client.completions.create(model=MODEL_NAME, prompt=SIMPLE_PROMPT)
​
    print("-" * 50)
    print(f"Completion results for {MODEL_NAME}:")
    print(completion)
    print("-" * 50)
​
​
def test_accuracy():
    """端到端精度测试:经聚合代理跑 gsm8k,与基线精度比对。"""
    run_simple_prompt()
​
    model_args = (
        f"model={MODEL_NAME},"
        f"base_url={BASE_URL}/completions,"
        f"num_concurrent={NUM_CONCURRENT},tokenized_requests=False"
    )
    results = lm_eval.simple_evaluate(
        model="local-completions",
        model_args=model_args,
        tasks=TASK,
    )
​
    measured_value = results["results"][TASK][FILTER]
    expected_value = EXPECTED_VALUES.get(MODEL_NAME)
​
    print(f"Measured accuracy value: {measured_value}\n")
    if expected_value is None:
        # 未登记的模型只输出精度、不做硬断言,便于 CI 之外手动验证。
        print(
            f"Warning: No expected value found for {MODEL_NAME}. "
            "Skipping accuracy check."
        )
        return
​
    # 容差判断:measured 需落在 expected ± RTOL 区间内,防止精度回归。
    assert (
        measured_value - RTOL < expected_value
        and measured_value + RTOL > expected_value
    ), f"Expected: {expected_value} | Measured: {measured_value}"
tests/v1/kv_connector/mooncake_integration/run_accuracy_test.sh test-coverage

编排核心:负责按环境变量拉起多组 prefill/decode vllm 实例、配置 Mooncake bootstrap 与 kv-transfer-config、等待就绪后启动聚合代理并触发 pytest;脚本内 RDMA 设备白名单注释记录了 b200-k8s 池的关键踩坑。

#!/bin/bash
set -xe
​
# ---- 配置来源:环境变量,全部带默认值 ----
NUM_PREFILL_INSTANCES=${NUM_PREFILL_INSTANCES:-1} # prefill 实例数
NUM_DECODE_INSTANCES=${NUM_DECODE_INSTANCES:-1} # decode 实例数
PREFILLER_TP_SIZE=${PREFILLER_TP_SIZE:-1}
DECODER_TP_SIZE=${DECODER_TP_SIZE:-1}
GPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.2}
BLOCK_SIZE=${BLOCK_SIZE:-128}# ---- Mooncake 设备发现 ----
# Mooncake 通过 bootstrap server 发现对端,不需要 side-channel 布局变量。
# MOONCAKE_DEVICE_NAME 是拓扑发现的 HCA 白名单;默认值针对 b200-k8s CI 池:
# 该池节点暴露 4 个 InfiniBand HCA 但只有一个 RoCE 口(mlx5_12),若全部 HCA
# 可见,P/D 两侧可能选中不同链路层,QP 握手无法到达 RTR。其他硬件需调整该值。
MOONCAKE_DEVICE_NAME=${MOONCAKE_DEVICE_NAME-mlx5_12}# producer(prefill 侧)与 consumer(decode 侧)共用配置模板,仅 kv_role 不同。
kv_config() {
  printf '{"kv_connector":"MooncakeConnector","kv_role":"%s",' "$1"
  printf '"kv_connector_extra_config":{"device_name":"%s"}}' "${MOONCAKE_DEVICE_NAME}"
}
KV_CONFIG_P=$(kv_config kv_producer)
KV_CONFIG_D=$(kv_config kv_consumer)
​
run_tests_for_model() {
  local model_name=$1  # 启动 prefill 实例:GPU 按实例序号轮转,TP > 1 时追加相邻 GPU。
  for i in $(seq 0 $((NUM_PREFILL_INSTANCES - 1))); do
    GPU_ID=$((i % $(get_num_gpus)))
    for ((j = 1; j < PREFILLER_TP_SIZE; j++)); do
      GPU_ID="${GPU_ID},$(((GPU_ID + j) % $(get_num_gpus)))"
    done    # 端口规划:prefill 8100+ / decode 8200+ / bootstrap 8998+,避免互相冲突。
    PORT=$((8100 + i))
    BOOTSTRAP_PORT=$((8998 + i))    eval "CUDA_VISIBLE_DEVICES=$GPU_ID \
      VLLM_MOONCAKE_BOOTSTRAP_PORT=$BOOTSTRAP_PORT \
      vllm serve $model_name \
      --port $PORT --enforce-eager \
      --block-size ${BLOCK_SIZE} \
      --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
      --tensor-parallel-size $PREFILLER_TP_SIZE \
      --kv-transfer-config '$KV_CONFIG_P' &"    PROXY_ARGS+=(--prefill "http://localhost:${PORT}" "$BOOTSTRAP_PORT")
  done  # decode 实例与 prefill 对称:端口 8200+、kv_role=kv_consumer。
  for i in $(seq 0 $((NUM_DECODE_INSTANCES - 1))); do
    GPU_ID=$(((i + NEXT_GPU + 1) % $(get_num_gpus)))
    for ((j = 1; j < DECODER_TP_SIZE; j++)); do
      GPU_ID="${GPU_ID},$(((GPU_ID + j) % $(get_num_gpus)))"
    done
    PORT=$((8200 + i))    eval "CUDA_VISIBLE_DEVICES=$GPU_ID \
      vllm serve $model_name \
      --port $PORT --enforce-eager \
      --block-size ${BLOCK_SIZE} \
      --gpu-memory-utilization $GPU_MEMORY_UTILIZATION \
      --tensor-parallel-size $DECODER_TP_SIZE \
      --kv-transfer-config '$KV_CONFIG_D' &"    PROXY_ARGS+=(--decode "http://localhost:${PORT}")
  done  # 所有 P/D 实例就绪后再启动聚合代理,统一暴露 8192 端口给下游测试。
  for PORT in "${PREFILL_PORTS[@]}" "${DECODE_PORTS[@]}"; do
    wait_for_server "$PORT" || return 1
  done  python3 "${GIT_ROOT}/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py" \
    --port 8192 "${PROXY_ARGS[@]}" &
  sleep 5  # 精度断言在 pytest 中执行:跑 gsm8k 并与模型基线做容差比较。
  TEST_MODEL=$model_name python3 -m pytest -s -x \
    "${GIT_ROOT}/tests/v1/kv_connector/mooncake_integration/test_accuracy.py"  # 每个模型跑完后清理残留 vllm serve 进程,避免端口与显存互相干扰。
  cleanup_instances
}
.buildkite/scripts/install-kv-connectors.sh core-logic

CI 基础设施关键改动:扩展元数据探测到 mooncake,并在 CUDA 13 镜像上自动把 CUDA 12 wheel 替换为 cuda13 变体,否则引擎启动即链接失败。

# ---- 收集已安装 KV connector 组件版本(nixl + mooncake)----
# 依 torch.version.cuda 判断 CUDA 主版本,用于后续 wheel 变体选择。
KV_METADATA=$(python3 - <<'PY'
import importlib.metadata as metadata
import torchcuda_version = torch.version.cuda
if cuda_version is None:
    raise SystemExit("torch.version.cuda is not set")try:
    mooncake_version = metadata.version("mooncake-transfer-engine")
except metadata.PackageNotFoundError:
    mooncake_version = ""print(cuda_version.split(".", 1)[0], metadata.version("nixl"), mooncake_version)
PY
)
read -r CUDA_MAJOR NIXL_VERSION MOONCAKE_VERSION <<< "${KV_METADATA}"
MOONCAKE_VERSION="${MOONCAKE_VERSION:-}"# ---- Mooncake CUDA 12/13 wheel 替换 ----
# PyPI 默认的 mooncake-transfer-engine wheel 基于 CUDA 12 构建,engine.so 链接
# libcudart.so.12,而 CUDA 13 运行镜像没有该库;mismatch 会以 ImportError
# 在引擎启动前暴露。因此 CUDA 13 镜像上先卸载默认 wheel,再装同版本 cuda13
# 变体——两者都提供 `mooncake` 包,必须先卸载避免包冲突。
if [ "${CUDA_MAJOR}" = "13" ] && [ -n "${MOONCAKE_VERSION}" ]; then
  uv pip uninstall --system mooncake-transfer-engine 2>/dev/null || true
  uv pip install --system "mooncake-transfer-engine-cuda13==${MOONCAKE_VERSION}"
fi

评论区精华

是否复用示例代理而非自建 toy proxy 设计

zhewenl 在 toy_proxy_server.py 上提问:"Why couldn't we just reuse the example proxy?",指出测试目录里新增的 toy 代理与 examples 里的 mooncake_connector_proxy.py 功能重复。

结论:NickLucche 承认 "I overfitted on nixl setup lol"(照搬了 nixl 套件的结构),最终版本删除了 toy_proxy_server.py,改为直接调用 examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py。 · 已解决

是否补充 SWA(滑动窗口注意力)模型覆盖 测试

zhewenl 建议补充 SWA 模型测试:"do you feel like we should also add some tests for SWA models, maybe https://huggingface.co/google/gemma-3-1b-it will be a good fit";ivanium 表示同意。

结论:NickLucche 回应 "this is a PR to get things started",当前优先打通基础设施(处理 linking errors 等),不同模型类型的配置可仿照 nixl 的做法在后续 PR 按模型补充;该建议作为后续工作记录,未在本 PR 实现。 · 待处理

风险与影响

风险集中在四方面:一是测试稳定性依赖节点分配——任务锁定 b200-k8s 池,RDMA 对 IPC_LOCK 与 host networking 是硬性依赖,作者在 Issue 评论中也提到节点分配导致的 RDMA 问题;MOONCAKE_DEVICE_NAME=mlx5_12 是强环境耦合,换池或换网卡都会直接失败。二是固定精度基线可能误报或漏报——Qwen/Qwen3-0.6B 的 0.41 基线来自本地 benchmark,模型或推理配置变化可能触发误报,而未登记模型只告警不断言又可能掩盖真实回归。三是CUDA wheel 替换可能静默失败——install-kv-connectors.sh 在 CUDA 13 镜像上先卸载再装 cuda13 变体,卸载用 || true 容错;若变体版本缺失,import canary 只能发现 mooncake.engine 导入问题,无法覆盖 libcudart 运行时链接层面的错配。四是超时窗口偏紧——Job 超时 10 分钟,但脚本要启动多实例、wait_for_server 单次最长 1200 秒并执行两轮配置扫描,节点启动慢或模型下载慢时容易整组超时。

对用户与系统无运行时影响(零源码主路径改动),影响集中在 CI 基础设施与 kv-connector 测试矩阵:Mooncake PD 首次获得端到端回归保护,且任务组可独立开关、仅在相关文件变更时触发,能节省 CI 开销。对团队而言,后续加模型、加 SWA 或更多 TP 配置只需在 sweep 脚本追加条目;该套脚本与 buildkite 配置也可作为其他 KV connector(如 nixl)之外新连接器的接入模板。影响程度中等,价值主要在工程保障层面。

RDMA 依赖特定 CI 节点类型 固定精度基线可能误报 CUDA wheel 替换可能静默失败 Job 超时窗口偏紧

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论