group: Disaggregated Mooncake
depends_on:
  - image-build
steps:
- label: ":nvidia: (B200) Distributed MooncakeConnector PD accuracy"
  key: distributed-mooncakeconnector-pd-accuracy-4-gpus
  timeout_in_minutes: 10
  working_dir: "/vllm-workspace/tests"
  # Runs on b200-k8s because it is the only NVIDIA pool that can carry mooncake
  # transfers: its pods are the only ones granted IPC_LOCK, which RDMA needs to
  # pin memory, and the only ones using host networking. Elsewhere topology
  # discovery finds no usable HCA and the tcp fallback then advertises the
  # docker bridge address, so every transfer fails to connect.
  device: b200-k8s
  num_devices: 4
  source_file_dependencies:
    - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
    - tests/v1/kv_connector/mooncake_integration/
  commands:
    - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
    - bash v1/kv_connector/mooncake_integration/config_sweep_accuracy_test.sh
  mirror:
    amd:
      label: ":amd: (MI355) Distributed MooncakeConnector HIP PD accuracy"
      dind: false
      device: mi355_4
      num_devices: 4
      timeout_in_minutes: 30
      working_dir: "/vllm-workspace/tests"
      depends_on:
        - image-build-amd
      source_file_dependencies:
        - vllm/distributed/kv_transfer/kv_connector/v1/mooncake/
        - tests/v1/kv_connector/mooncake_integration/
        - tests/v1/kv_connector/rocm_pd_accuracy_utils.py
        - tests/utils.py
        - examples/disaggregated/mooncake_connector/
        - vllm/v1/attention/backends/rocm_aiter_unified_attn.py
        - vllm/platforms/rocm.py
        - requirements/common.txt
        - requirements/kv_connectors_rocm.txt
      commands:
        - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
        # Same full accuracy/TP sweep, using same-host HIP GPU transport.
        - pytest -v -s v1/kv_connector/mooncake_integration/test_accuracy_rocm.py

- label: ":nvidia: (L4) Mooncake EC TCP E2E"
  key: mooncake-ec-tcp-e2e-2-gpus
  timeout_in_minutes: 30
  working_dir: "/vllm-workspace"
  device: l4
  num_devices: 2
  source_file_dependencies:
    - vllm/distributed/ec_transfer/
    - vllm/config/ec_transfer.py
    - vllm/config/multimodal.py
    - vllm/config/vllm.py
    - vllm/multimodal/
    - vllm/v1/core/encoder_cache_manager.py
    - vllm/v1/core/sched/
    - vllm/v1/engine/
    - vllm/v1/worker/ec_connector_model_runner_mixin.py
    - vllm/v1/worker/gpu_model_runner.py
    - vllm/v1/worker/gpu_worker.py
    - vllm/v1/worker/gpu/ec_connector.py
    - vllm/v1/worker/gpu/model_runner.py
    - vllm/v1/worker/mm_encoder_model_runner.py
    - vllm/v1/worker/encoder_cudagraph.py
    - vllm/v1/worker/encoder_cudagraph_defs.py
    - vllm/v1/attention/ops/vit_attn_wrappers.py
    - vllm/v1/worker/gpu/model_states/
    - vllm/v1/worker/gpu/mm/
    - tests/v1/cudagraph/test_encoder_cudagraph.py
    - tests/models/multimodal/generation/test_vit_cudagraph.py
    - examples/disaggregated/disaggregated_encoder/
    - tests/v1/ec_connector/
    - requirements/kv_connectors.txt
  env:
    PYTHON_BIN: "/vllm-workspace/.venv/bin/python"
    PYTHONUNBUFFERED: "1"
    VLLM_HOST_IP: "127.0.0.1"
    VLLM_USE_V2_MODEL_RUNNER: "1"
    E_CUDAGRAPH_MM_ENCODER: "1"
    # FA padding NaNs are tracked separately in #57136.
    E_MM_ENCODER_ATTN_BACKEND: "TRITON_ATTN"
    MOONCAKE_EC_PROTOCOL: "tcp"
    USE_MM_PROMPTS: "1"
    SKIP_BASELINE: "0"
    CONCURRENCY: "3"
    REPEAT: "2"
    LOG_PATH: "/tmp/mooncake-ec-e2e"
    BASELINE_FILE: "/tmp/mooncake-ec-e2e/baseline.json"
  commands:
    - uv venv --system-site-packages --python 3.12 .venv
    - |
      uv pip install --python .venv/bin/python "$(.venv/bin/python - <<'PY'
      from pathlib import Path
      import torch

      requirement = next(
          line for line in Path("requirements/kv_connectors.txt").read_text().splitlines()
          if line.startswith("mooncake-transfer-engine ")
      )
      if torch.version.cuda.split(".")[0] == "13":
          requirement = requirement.replace("mooncake-transfer-engine", "mooncake-transfer-engine-cuda13")
      print(requirement)
      PY
      )"
    - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
  mirror:
    amd:
      label: ":amd: (MI300) Mooncake EC TCP E2E"
      dind: false
      device: mi300_2
      num_devices: 2
      timeout_in_minutes: 60
      working_dir: /vllm-workspace
      depends_on:
      - image-build-amd
      source_file_dependencies:
      - vllm/distributed/ec_transfer/
      - vllm/config/ec_transfer.py
      - vllm/config/multimodal.py
      - vllm/config/vllm.py
      - vllm/multimodal/
      - vllm/v1/core/encoder_cache_manager.py
      - vllm/v1/core/sched/
      - vllm/v1/engine/
      - vllm/v1/worker/ec_connector_model_runner_mixin.py
      - vllm/v1/worker/gpu_model_runner.py
      - vllm/v1/worker/gpu_worker.py
      - vllm/v1/worker/gpu/ec_connector.py
      - vllm/v1/worker/gpu/model_runner.py
      - vllm/v1/worker/mm_encoder_model_runner.py
      - vllm/v1/worker/encoder_cudagraph.py
      - vllm/v1/worker/encoder_cudagraph_defs.py
      - vllm/v1/attention/ops/vit_attn_wrappers.py
      - vllm/v1/worker/gpu/model_states/
      - vllm/v1/worker/gpu/mm/
      - tests/v1/cudagraph/test_encoder_cudagraph.py
      - tests/models/multimodal/generation/test_vit_cudagraph.py
      - examples/disaggregated/disaggregated_encoder/
      - tests/v1/ec_connector/
      - requirements/kv_connectors.txt
      - vllm/platforms/rocm.py
      - requirements/rocm.txt
      env:
        VLLM_WORKER_MULTIPROC_METHOD: "spawn"
        PYTHON_BIN: "python"
      commands:
      - bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
