support hybrid attention for vllm connector && add integration test #1
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| name: test-vllm-e2e | |
| permissions: | |
| contents: read | |
| on: | |
| pull_request: | |
| branches: ["main"] | |
| workflow_dispatch: | |
| inputs: | |
| runs-on: | |
| description: "GPU runner label (needs >= 2 GPUs, e.g. A10)" | |
| type: string | |
| default: "gpu-a10-x2" | |
| concurrency: | |
| group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} | |
| cancel-in-progress: ${{ github.event_name == 'pull_request' }} | |
| jobs: | |
| vllm-e2e: | |
| name: vllm-e2e (${{ matrix.kind }}) | |
| # GPU runners are self-hosted; the default label can be overridden per-repo | |
| # via the VLLM_E2E_RUNS_ON variable or the workflow_dispatch input. | |
| runs-on: ${{ inputs.runs-on || vars.VLLM_E2E_RUNS_ON || 'gpu-a10-x2' }} | |
| timeout-minutes: 180 | |
| strategy: | |
| fail-fast: false | |
| matrix: | |
| include: | |
| # Full-attention model: single FullAttentionSpec kv cache group. | |
| - kind: full-attention | |
| model_var: VLLM_E2E_MODEL_FULL_ATTN | |
| # Hybrid model: MambaSpec groups + FullAttentionSpec group | |
| # (mamba_cache_mode="align"). | |
| - kind: hybrid-attention | |
| model_var: VLLM_E2E_MODEL_HYBRID | |
| steps: | |
| - uses: actions/checkout@v4 | |
| - name: check_gpus | |
| run: | | |
| nvidia-smi --query-gpu=index,name,memory.total --format=csv,noheader | |
| GPU_COUNT=$(nvidia-smi --query-gpu=name --format=csv,noheader | wc -l) | |
| if [ "$GPU_COUNT" -lt 2 ]; then | |
| echo "::error::This test needs at least 2 GPUs, found $GPU_COUNT" | |
| exit 1 | |
| fi | |
| - name: resolve_runner_config | |
| # Self-hosted runner provisioning: a vLLM (>= 0.26.0) venv and the model | |
| # checkouts are pre-cached on the runner and exposed via repo variables: | |
| # VLLM_E2E_PYTHON venv python with vllm installed | |
| # VLLM_E2E_MODEL_FULL_ATTN e.g. a local Qwen2.5-7B-Instruct checkout | |
| # VLLM_E2E_MODEL_HYBRID e.g. a local Qwen3.5-4B checkout | |
| env: | |
| VENV_PY: ${{ vars.VLLM_E2E_PYTHON }} | |
| MODEL: ${{ vars[matrix.model_var] }} | |
| run: | | |
| if [ ! -x "$VENV_PY" ]; then | |
| echo "::error::VLLM_E2E_PYTHON ($VENV_PY) not found; provision the runner" | |
| exit 1 | |
| fi | |
| "$VENV_PY" -c 'import vllm; v = vllm.__version__; print("vllm", v)' | |
| if [ ! -f "$MODEL/config.json" ]; then | |
| echo "::error::model not found at $MODEL; provision the runner" | |
| exit 1 | |
| fi | |
| echo "KVCM_E2E_PYTHON=$VENV_PY" >> "$GITHUB_ENV" | |
| echo "KVCM_E2E_MODEL=$MODEL" >> "$GITHUB_ENV" | |
| - name: build_binaries | |
| run: | | |
| set -x | |
| bazelisk build //kv_cache_manager:kv_cache_manager_bin \ | |
| //kv_cache_manager/client/pybind:kvcm_py_client_lib_wheel \ | |
| //kv_cache_manager/py_connector/vllm:kvcm_vllm_connector_wheel \ | |
| --per_file_copt='external/jsoncpp_git/.*@-Wno-error' | |
| - name: install_wheels | |
| # Bazel wheel filenames contain unstamped {STABLE_*} template variables; | |
| # read the real version from the wheel METADATA and rename before install. | |
| run: | | |
| set -x | |
| mkdir -p /tmp/kvcm_whl && rm -f /tmp/kvcm_whl/*.whl | |
| for whl in bazel-bin/kv_cache_manager/client/pybind/kvcm_py_client-*.whl \ | |
| bazel-bin/kv_cache_manager/py_connector/vllm/kvcm_vllm_connector-*.whl; do | |
| pkg=$(basename "$whl" | sed 's/-{STABLE.*//') | |
| ver=$(unzip -p "$whl" "*.dist-info/METADATA" | awk '/^Version:/{print $2; exit}') | |
| cp "$whl" "/tmp/kvcm_whl/${pkg}-${ver}-cp312-cp312-manylinux_2_32_x86_64.whl" | |
| done | |
| "$KVCM_E2E_PYTHON" -m pip install --no-deps --force-reinstall /tmp/kvcm_whl/*.whl || \ | |
| uv pip install --python "$KVCM_E2E_PYTHON" --no-deps --force-reinstall /tmp/kvcm_whl/*.whl | |
| - name: run_e2e_tests | |
| run: | | |
| set -x | |
| bazelisk test //integration_test/vllm_e2e/... \ | |
| --cache_test_results=no --test_output=errors \ | |
| --test_env=KVCM_E2E_PYTHON="$KVCM_E2E_PYTHON" \ | |
| --test_env=KVCM_E2E_MODEL="$KVCM_E2E_MODEL" \ | |
| --per_file_copt='external/jsoncpp_git/.*@-Wno-error' | |
| - name: upload_logs | |
| if: failure() | |
| uses: actions/upload-artifact@v6 | |
| with: | |
| name: vllm-e2e-logs-${{ matrix.kind }} | |
| path: | | |
| /tmp/kvcm_vllm_e2e/**/*.stdout | |
| /tmp/kvcm_vllm_e2e/**/*.stderr | |
| bazel-out/*-opt/testlogs/integration_test/vllm_e2e/** | |
| if-no-files-found: ignore |