iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Mini-VLLM] Remove orphan subsystems: snapshot, CPU workers, dead decoders

Browse Source 
- Delete the snapshot package, CLI command, and its launch tests
- Delete CPU workers, cpu_communicator, and the hw_agnostic layer package
- Delete the EC CPU connector, keeping the EC base glue intact
- Delete dead composite/triton_flashinfer backends and third_party/flashmla
- Delete dead video decoders, video-bound tests, and orphan utils
- Remove the deep_gemm_mega_moe option and dead allreduce fusion pair

Co-authored-by: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
MOVZX
2026-10-19 09:27:06 +07:00
co-authored by Hinata
parent 2b544b7ba0
commit bd70a61c20
88 changed files with 4 additions and 20247 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/entrypoints/launchers/test_launch_cli.py
-1436
View File
File diff suppressed because it is too large Load Diff
tests/model_executor/test_jit_warmup_cutedsl_launcher.py
-111
View File
@@ -1,111 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from dataclasses import dataclass
from typing import Any
from vllm.model_executor.warmup import jit_warmup_cutedsl_helper
from vllm.model_executor.warmup.jit_warmup import kernel_launcher
from vllm.model_executor.warmup.jit_warmup_cutedsl_helper import (
CuTeDSLLaunchSpec,
VllmCuTeDSLJitKernel,
)
class _TestCuTeDSLKernel(VllmCuTeDSLJitKernel["_TestCuTeDSLKernel.CompileKey"]):
@dataclass(frozen=True)
class CompileKey:
variant: int
@staticmethod
def kernel(compile_key: CompileKey) -> tuple[str, int]:
return "entry", compile_key.variant
def dispatch(self, *, variant: int) -> CompileKey:
return self.CompileKey(variant=variant)
def get_warmup_keys(self) -> list[CompileKey]:
return self._trace_dispatch(self.dispatch)(variant=(1, 2))
def warmup_inputs(self, compile_key: CompileKey) -> tuple[Any, ...]:
return (f"fake-{compile_key.variant}",)
@kernel_launcher
def __call__(
self,
payload: str,
*,
variant: int,
) -> CuTeDSLLaunchSpec[CompileKey]:
compile_key = self.dispatch(variant=variant)
return compile_key, (payload,)
def test_cutedsl_launcher_reuses_compiled_executor(monkeypatch) -> None:
compile_calls: list[tuple[Any, tuple[Any, ...]]] = []
launch_calls: list[tuple[int, tuple[Any, ...]]] = []
def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any:
compile_calls.append((entry, args))
variant = entry[1]
def executor(*runtime_args: Any) -> tuple[int, tuple[Any, ...]]:
launch_calls.append((variant, runtime_args))
return variant, runtime_args
return executor
monkeypatch.setattr(
jit_warmup_cutedsl_helper,
"compile_cutedsl",
compile_cutedsl,
)
kernel = _TestCuTeDSLKernel()
compile_key = kernel.CompileKey(variant=1)
kernel.compile(compile_key)
kernel.compile(compile_key)
assert compile_calls == [(("entry", 1), ("fake-1",))]
assert kernel("runtime", variant=1) == (1, ("runtime",))
assert kernel("other", variant=2) == (2, ("other",))
assert compile_calls == [
(("entry", 1), ("fake-1",)),
(("entry", 2), ("fake-2",)),
]
assert launch_calls == [(1, ("runtime",)), (2, ("other",))]
def test_cutedsl_launcher_runs_post_launch_finalizer(monkeypatch) -> None:
events: list[str] = []
def compile_cutedsl(entry: tuple[str, int], *args: Any) -> Any:
def executor(*runtime_args: Any) -> None:
events.append("launch")
return executor
class FinalizedKernel(_TestCuTeDSLKernel):
@kernel_launcher
def __call__(
self,
payload: str,
*,
variant: int,
) -> CuTeDSLLaunchSpec[_TestCuTeDSLKernel.CompileKey]:
compile_key = self.dispatch(variant=variant)
def finalize() -> str:
events.append("finalize")
return "output"
return compile_key, (payload,), "unused", finalize
monkeypatch.setattr(
jit_warmup_cutedsl_helper,
"compile_cutedsl",
compile_cutedsl,
)
assert FinalizedKernel()("runtime", variant=1) == "output"
assert events == ["launch", "finalize"]
tests/models/transformers/test_layer_registry_e2e.py
-113
View File
@@ -1,113 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end check of the Transformers backend's hw-agnostic layers.
Serves a tiny Llama through `vllm_runner` with and without
`VLLM_USE_HW_AGNOSTIC` and checks the layer providers and logprobs. Needs a
real engine, so it runs in the GPU job. The CPU-only unit tests for the
resolution logic (`layers._resolve`) are in `test_layer_registry.py`.
"""
import pytest
import torch
@pytest.fixture(scope="module")
def tiny_llama_path(tmp_path_factory):
"""A randomly-initialized microscopic Llama saved to disk (with an ungated
tokenizer) so vLLM can load it like any local checkpoint."""
from transformers import AutoTokenizer, LlamaConfig, LlamaForCausalLM
tokenizer = AutoTokenizer.from_pretrained("hf-internal-testing/llama-tokenizer")
config = LlamaConfig(
vocab_size=tokenizer.vocab_size,
hidden_size=128,
intermediate_size=256,
num_hidden_layers=2,
num_attention_heads=4,
rms_norm_eps=1e-6,
hidden_act="silu",
)
torch.manual_seed(0)
model = LlamaForCausalLM(config)
path = tmp_path_factory.mktemp("tiny_llama")
model.save_pretrained(path)
tokenizer.save_pretrained(path)
return str(path)
# Registered names of the layers the backend can
# currently route to hw-agnostic implementations.
_COVERED_LAYERS = ("rms_norm", "silu_and_mul")
def _layer_providers(model) -> dict[str, str]:
"""Map each covered layer type present in the model to the provider its
implementation came from (``hw_agnostic`` or ``vllm``).
"""
def provider_of(module) -> str | None:
for cls in type(module).__mro__:
if "hw_agnostic.layers" in cls.__module__:
return "hw_agnostic"
if ".model_executor.layers." in cls.__module__:
return "vllm"
return None
providers: dict[str, str] = {}
for module in model.modules():
name = getattr(module, "name", None)
if (
name in _COVERED_LAYERS
and name not in providers
and (prov := provider_of(module)) is not None
):
providers[name] = prov
return providers
def _serve(vllm_runner, model_path, prompts):
"""Serve the model through the backend; return (layer_providers, logprobs)."""
with vllm_runner(
model_path,
model_impl="transformers",
max_model_len=64,
enforce_eager=True,
gpu_memory_utilization=0.3,
) as runner:
assert runner.llm.llm_engine.model_config.using_transformers_backend()
providers = runner.apply_model(_layer_providers)[0]
outputs = runner.generate_greedy_logprobs(
prompts, max_tokens=32, num_logprobs=5
)
return providers, outputs
def test_hw_agnostic_matches_vllm_end_to_end(monkeypatch, vllm_runner, tiny_llama_path):
"""Serving the tiny model with hw-agnostic layers matches the vLLM baseline."""
# spawn: `layers` resolves its symbols at import time, so the worker must
# re-import it with VLLM_USE_HW_AGNOSTIC already set (fork would inherit
# the parent's already-resolved module).
monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
# apply_model pickles the introspection function.
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
from ..utils import check_logprobs_close
prompts = ["The capital of France is", "vLLM is"]
monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "0")
vllm_providers, vllm_outputs = _serve(vllm_runner, tiny_llama_path, prompts)
# Both replaceable layers present in a Llama block must be vLLM's here.
assert vllm_providers == {"rms_norm": "vllm", "silu_and_mul": "vllm"}
monkeypatch.setenv("VLLM_USE_HW_AGNOSTIC", "1")
hw_providers, hw_outputs = _serve(vllm_runner, tiny_llama_path, prompts)
assert hw_providers == {"rms_norm": "hw_agnostic", "silu_and_mul": "hw_agnostic"}
check_logprobs_close(
outputs_0_lst=vllm_outputs,
outputs_1_lst=hw_outputs,
name_0="vllm",
name_1="hw_agnostic",
)
tests/v1/ec_connector/__init__.py
View File
tests/v1/ec_connector/integration/README.md
-233
View File
@@ -1,233 +0,0 @@
# EPD Correctness Test
This test verifies that EPD (Encoder-Prefill-Decode) disaggregation produces identical outputs to a baseline single instance.
## What It Tests
- **Baseline**: Single vLLM instance serving a multimodal model
- **EPD (1E+1PD)**: 1 Encoder + 1 Prefill-Decode instance
- **Baseline (1P+1D)**: 1 Prefill + 1 Decode instance
- **EPD (1E+1P+1D)**: 1 Encoder + 1 Prefill + 1 Decode instance
The test ensures that disaggregated encoding produces **identical** outputs to the baseline.
Note that currently PD disaggregation set up may give slightly different results from a single instance. Therefore, we need the result from 1P+1D as the baseline for 1E+1P+1D
Please refer to [Disaggregated Encoder Feature](../../../../docs/features/disagg_encoder.md) for the detailed explanation for the EPD features.
## Files
- `run_epd_correctness_test.sh` - Main test script (starts all instances and runs tests)
- `test_epd_correctness.py` - Python test script (compares outputs)
## Usage
### Multimodal Prompts (Default)
```bash
cd vllm
./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
```
This runs the test with actual multimodal (image) prompts.
### Text-Only Prompts
```bash
cd vllm
USE_MM_PROMPTS=0 ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
```
This runs a quick test with text-only prompts to verify the setup works.
### Custom Configuration
```bash
# Use specific GPUs
GPU_E=0 GPU_PD=1 GPU_P=1 GPU_D=2 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
# Use specific ports
ENDPOINT_PORT=10001 bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
# Use specific model
MODEL="Qwen/Qwen2.5-VL-3B-Instruct" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
# Use specific storage path
EC_SHARED_STORAGE_PATH="/tmp/my_ec_cache" bash ./tests/v1/ec_connector/integration/run_epd_correctness_test.sh
```
## How It Works
### NIXL EC failure isolation (1E + 1PD)
```bash
PATH="$PWD/.venv/bin:$PATH" CUDA_VISIBLE_DEVICES=0,1 \
.venv/bin/python -m pytest \
tests/v1/ec_connector/integration/test_nixl_failure.py -v -s
```
Requires two CUDA GPUs and NIXL. `MODEL` overrides the default
`Qwen/Qwen2.5-VL-3B-Instruct`. The test starts real E/PD servers with CUDA
graphs enabled and uses the proxy's rewrite helper. Only the failing request's
control endpoint is replaced: a ZMQ peer returns `NACK_MISSING` and verifies
that the consumer actually requested the encoding.
Metadata-only requests must return a request-level error; requests retaining
the image must succeed through local fallback. After each case, a fresh
metadata-only image request must succeed through the real NIXL transfer path.
No server restart, proxy retry, cache hit, or timeout race can mask the result.
This is a correctness test, not a throughput benchmark.
### Mooncake EC (1E + 1PD)
```bash
PYTHON_BIN="$PWD/.venv/bin/python" \
bash tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
```
Requires two GPUs and Mooncake TransferEngine. TCP is the default transport;
no RDMA-capable network hardware is required.
The script sets `MC_FORCE_TCP=1` in TCP mode so Mooncake cannot auto-select RDMA.
The baseline runs first on GPU 0, followed by E on GPU 0 and PD on GPU 1.
`MOONCAKE_EC_PROTOCOL=rdma` selects RDMA instead; host-specific transport
environment variables should be set by the caller.
Three black-box cases check fixed short answers and compare with the baseline:
- One image: read the STOP sign.
- Two different images (including a local file): identify flowers and birds.
- The same image twice: read both STOP signs.
By default, all three requests run concurrently for two rounds, exercising
shared hashes across requests and reuse after completion. Every response is
compared, not just the final round. Set `CONCURRENCY` and `REPEAT` to override.
Prefix caching is disabled and CUDA graphs remain enabled. This is a small
correctness suite, not a performance benchmark or failure-injection suite.
`LOG_PATH` and `BASELINE_FILE` select the log directory and reference output.
`SKIP_BASELINE=1` reuses a reference generated with the same model and test
configuration. `USE_MM_PROMPTS=0` only checks text routing, not EC transfer.
Buildkite runs this script as `mooncake-ec-tcp-e2e-2-gpus` on two L4 GPUs.
The job is defined in `.buildkite/test_areas/disaggregated_mooncake.yaml` and
selected for changes to EC, multimodal processing, scheduler/model-runner
integration, the proxy, or these tests (subject to normal PR CI approval).
It reuses the CI image's Python packages through a system-site-packages venv
and installs the CUDA-compatible Mooncake wheel. It uses loopback networking
and TCP only, without RDMA devices or peer-memory setup.
The job fails on startup errors, request errors, or answer mismatches. On
failure, the script prints the last 100 lines of each server/proxy log to
the CI job log; full files remain under `LOG_PATH` while the container exists.
### Step 1: Baseline
1. Start single vLLM instance on GPU
2. Run test prompts (multimodal or text-only)
3. Save outputs to `.vllm_epd_baseline.txt`
4. Shutdown instance
### Step 2: EPD (1E + 1PD)
1. Clear encoder cache storage
2. Start instances and proxy
3. Run same test prompts
4. Assert outputs match baseline exactly
5. Shutdown instances
### Step 3: EPD (1E + 1P + 1D)
1. Clear encoder cache storage
2. Start instances and proxy
3. Run same test prompts
4. Assert outputs match baseline exactly
5. Shutdown instances
## Test Scenarios
### Multimodal Prompts (--use_mm_prompts)
Tests encoder cache transfer:
- Single image query
- Multiple images in one request
- Mixed image and text
- Image with detailed questions
### Text-Only Prompts (default)
Quick sanity check:
- Simple text queries
- Text-only explanations
- Verifies proxy routing works
## Expected Behavior
### ✅ Test Passes When
- All disagg outputs match baseline outputs exactly
- No errors during instance startup
- Encoder cache is properly saved and loaded
- Proxy correctly routes requests
### ❌ Test Fails When
- Outputs differ between baseline and disagg
- Server startup fails
- Encoder cache not found (should fall back to local execution)
- Proxy routing errors
## Notes
- The test uses deterministic generation (`temperature=0.0`, `seed=42`)
- Encoder cache should enable exact output reproduction
- Test cleans up all instances and cache files after completion
- Safe to run multiple times (idempotent)
- We setup the PD disagg part with NixlConnector. Please read details about EPD in `examples/disaggregated/disaggregated_encoder/README.md`
## Requirements
- Multiple GPUs (3 for 1E+1P+1D, 2 for 1E+1PD, 1 for baseline)
- 1E+1P+1D is runnable with 2 GPU by assign E and P on the same GPU now.
- Multimodal model (e.g., Qwen2.5-VL-3B-Instruct)
- Internet access (for accessing vllm test images)
## Debugging
### Check Logs
Logs and baseline output are saved in `/tmp/` by default.
Can be customized by changing the environment variables.
### Check Encoder Cache
```bash
# Verify cache files are created
ls -la $EC_SHARED_STORAGE_PATH/
# Should see directories with mm_hash names
# Each containing encoder_cache.safetensors
```
### Manual Testing
Run individual components:
```bash
# Baseline only
python test_epd_correctness.py \
--service_url http://localhost:8000 \
--model_name Qwen/Qwen2.5-VL-3B-Instruct \
--mode baseline \
--baseline_file test_output.txt \
--use_mm_prompts
# Disagg only (requires baseline output file!)
python test_epd_correctness.py \
--service_url http://localhost:8000 \
--model_name Qwen/Qwen2.5-VL-3B-Instruct \
--mode disagg \
--baseline_file test_output.txt \
--use_mm_prompts
```
tests/v1/ec_connector/integration/hato.jpg
BIN
View File
Binary file not shown.
Side by Side

Before

Width:  |  Height:  |  Size: 821 KiB

tests/v1/ec_connector/integration/run_epd_correctness_test.sh
-505
View File
@@ -1,505 +0,0 @@
#!/bin/bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# EPD (Encoder-Prefill-Decode) Correctness Test
#
# This script tests that EPD disaggregation produces the same outputs as baseline.
# It runs:
# 1. Baseline: Single vLLM instance
# 2. EPD: 1E + 1PD setup
# 3. Baseline for (E + P + D): 1P + 1D vLLM instances disagg
# 4. EPD: 1E + 1P + 1D setup
# For GPU usage
# set -xe
# Resolve the repository root from the script location instead of `.git`.
SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd -P)"
GIT_ROOT="${GIT_ROOT:-$(cd -- "${SCRIPT_DIR}/../../../.." && pwd -P)}"
# Model to test
MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
MAX_MODEL_LEN="${MAX_MODEL_LEN:-10240}"
GPU_MEMORY_UTILIZATION="${GPU_MEMORY_UTILIZATION:-0.7}"
MAX_NUM_SEQS="${MAX_NUM_SEQS:-128}"
# Set 1 to use multimodal prompts; else to use text-only
USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}"
USE_TWO_IMAGE_PROMPT="${USE_TWO_IMAGE_PROMPT:-1}"
TEST_FLAGS=()
if [[ "$USE_MM_PROMPTS" == "1" ]]; then
TEST_FLAGS+=(--use_mm_prompts)
fi
if [[ "$USE_TWO_IMAGE_PROMPT" != "1" ]]; then
TEST_FLAGS+=(--skip_two_image_prompt)
fi
# GPU configuration
GPU_E="${GPU_E:-0}"
GPU_P="${GPU_P:-1}"
GPU_D="${GPU_D:-2}"
GPU_SINGLE="${GPU_SINGLE:-$GPU_P}"
GPU_PD="${GPU_PD:-$GPU_P}"
# Device platform and affinity environment variable
DEVICE_PLATFORM="${DEVICE_PLATFORM:-cuda}"
if [[ -z "${DEVICE_AFFINITY_ENV:-}" ]]; then
if [[ "${DEVICE_PLATFORM,,}" == "xpu" ]]; then
DEVICE_AFFINITY_ENV="ZE_AFFINITY_MASK"
else
DEVICE_AFFINITY_ENV="CUDA_VISIBLE_DEVICES"
fi
fi
# Port
ENCODE_PORT="${ENCODE_PORT:-19534}"
PREFILL_PORT="${PREFILL_PORT:-19535}"
DECODE_PORT="${DECODE_PORT:-19536}"
PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19537}"
ENDPOINT_PORT="${ENDPOINT_PORT:-10001}"
# Storage path for encoder cache
EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache_test}"
TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-600}"
# Output file for baseline comparison and logs
LOG_PATH="${LOG_PATH:-/tmp}"
BASELINE_FILE="${BASELINE_FILE:-/tmp/vllm_baseline.txt}"
BASELINE_PD_FILE="${BASELINE_PD_FILE:-/tmp/vllm_epd_baseline.txt}"
mkdir -p "$LOG_PATH"
# Trap the SIGINT signal (triggered by Ctrl+C)
trap 'kill $(jobs -pr)' SIGINT SIGTERM EXIT
# Wait for server to be ready
wait_for_server() {
local port=$1
timeout "$TIMEOUT_SECONDS" bash -c "
until curl -s localhost:${port}/v1/chat/completions > /dev/null; do
sleep 1
done" && return 0 || return 1
}
# Cleanup function
cleanup_instances() {
echo "Cleaning up any running vLLM instances..."
pkill -f "vllm serve" || true
pkill -f "disagg_epd_proxy.py" || true
sleep 2
}
# Function to run baseline (single instance)
run_baseline() {
echo "================================"
echo "Running BASELINE (single instance)"
echo "================================"
cleanup_instances
rm -rf "$EC_SHARED_STORAGE_PATH"
local PORT=$ENDPOINT_PORT
# Start baseline instance
echo "Starting baseline instance on GPU $GPU_SINGLE, port $PORT"
env "$DEVICE_AFFINITY_ENV=$GPU_SINGLE" vllm serve "$MODEL" \
--port "$PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--enforce-eager \
--gpu-memory-utilization 0.9 \
--max-num-seqs "$MAX_NUM_SEQS" \
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
> "$LOG_PATH"/baseline.log 2>&1 &
local BASELINE_PID=$!
# Wait for baseline to start
echo "Waiting for baseline instance to start..."
wait_for_server "$PORT"
curl http://127.0.0.1:"$PORT"/v1/models
echo ""
# Run test in baseline mode
echo "Running baseline..."
python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
--service_url "http://localhost:$PORT" \
--model_name "$MODEL" \
--mode baseline \
--baseline_file "$BASELINE_FILE" \
"${TEST_FLAGS[@]}"
# Cleanup baseline
echo "Stopping baseline instance..."
kill $BASELINE_PID 2>/dev/null || true
sleep 2
cleanup_instances
}
# Function to run EPD with 1E + 1PD
run_epd_1e_1pd() {
echo "================================"
echo "Running EPD (1E + 1PD)"
echo "================================"
cleanup_instances
rm -rf "$EC_SHARED_STORAGE_PATH"
mkdir -p "$EC_SHARED_STORAGE_PATH"
local ENCODE_PORT=$ENCODE_PORT
local PREFILL_DECODE_PORT=$PREFILL_DECODE_PORT
local PROXY_PORT=$ENDPOINT_PORT
declare -a PIDS=()
# Start encoder instance
echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT"
env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \
--port "$ENCODE_PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--enforce-eager \
--gpu-memory-utilization 0.01 \
--enable-request-id-headers \
--no-enable-prefix-caching \
--max-num-batched-tokens 114688 \
--max-num-seqs "$MAX_NUM_SEQS" \
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_producer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}' \
> "$LOG_PATH"/1e1pd_encoder.log 2>&1 &
PIDS+=($!)
# Start prefill+decode instance
echo "Starting PD instance on GPU $GPU_PD, port $PREFILL_DECODE_PORT"
env "$DEVICE_AFFINITY_ENV=$GPU_PD" vllm serve "$MODEL" \
--port "$PREFILL_DECODE_PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--enforce-eager \
--enable-mm-embeds \
--gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
--enable-request-id-headers \
--max-num-seqs "$MAX_NUM_SEQS" \
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_consumer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}' \
> "$LOG_PATH"/1e1pd_pd.log 2>&1 &
PIDS+=($!)
# Wait for instances to start
echo "Waiting for encoder instance..."
wait_for_server "$ENCODE_PORT"
echo "Waiting for PD instance..."
wait_for_server "$PREFILL_DECODE_PORT"
# Start proxy
echo "Starting EPD proxy on port $PROXY_PORT"
python "${GIT_ROOT}/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" \
--host "0.0.0.0" \
--port "$PROXY_PORT" \
--encode-servers-urls "http://localhost:$ENCODE_PORT" \
--prefill-servers-urls "disable" \
--decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
> "$LOG_PATH"/1e1pd_proxy.log 2>&1 &
PIDS+=($!)
# Wait for proxy
echo "Waiting for proxy..."
wait_for_server "$PROXY_PORT"
curl http://127.0.0.1:"$PROXY_PORT"/v1/models
curl http://127.0.0.1:"$PROXY_PORT"/health
echo ""
echo "All EPD (1E+1PD) services are up!"
# Run test in disagg mode
echo "Running EPD (1E+1PD) correctness test..."
python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
--service_url "http://localhost:$PROXY_PORT" \
--model_name "$MODEL" \
--mode disagg \
--baseline_file "$BASELINE_FILE" \
"${TEST_FLAGS[@]}"
# Cleanup
echo "✓✓ 1E+1PD Correctness Test finished"
echo "Stopping EPD (1E+1PD) instances..."
for pid in "${PIDS[@]}"; do
kill "$pid" 2>/dev/null || true
done
sleep 2
cleanup_instances
}
# Function to run baseline for 1E + 1P + 1D (PD disagg)
run_baseline_1p_1d() {
echo "================================"
echo "Running PD BASELINE (1P + 1D)"
echo "================================"
cleanup_instances
rm -rf "$EC_SHARED_STORAGE_PATH"
mkdir -p "$EC_SHARED_STORAGE_PATH"
local PREFILL_PORT=$PREFILL_PORT
local DECODE_PORT=$DECODE_PORT
local PROXY_PORT=$ENDPOINT_PORT
declare -a PIDS=()
# Start prefill instance
echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT"
env "$DEVICE_AFFINITY_ENV=$GPU_P" \
VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
vllm serve "$MODEL" \
--port "$PREFILL_PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--enforce-eager \
--gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
--enable-request-id-headers \
--max-num-seqs "$MAX_NUM_SEQS" \
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
--kv-transfer-config '{
"kv_connector": "NixlConnector",
"kv_role": "kv_producer"
}' \
> "$LOG_PATH"/1p1d_prefill.log 2>&1 &
PIDS+=($!)
# Start decode instance
echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT"
env "$DEVICE_AFFINITY_ENV=$GPU_D" \
VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
vllm serve "$MODEL" \
--port "$DECODE_PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--enforce-eager \
--gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
--enable-request-id-headers \
--max-num-seqs "$MAX_NUM_SEQS" \
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
--kv-transfer-config '{
"kv_connector": "NixlConnector",
"kv_role": "kv_consumer"
}' \
> "$LOG_PATH"/1p1d_decode.log 2>&1 &
PIDS+=($!)
# Wait for instances to start
echo "Waiting for prefill instance..."
wait_for_server "$PREFILL_PORT"
echo "Waiting for decode instance..."
wait_for_server "$DECODE_PORT"
# Start proxy
echo "Starting EPD proxy on port $PROXY_PORT"
python "${GIT_ROOT}/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py" \
--host "0.0.0.0" \
--port "$PROXY_PORT" \
--prefiller-ports "$PREFILL_PORT" \
--decoder-ports "$DECODE_PORT" \
> "$LOG_PATH"/1p1d_proxy.log 2>&1 &
PIDS+=($!)
# Wait for proxy
echo "Waiting for proxy..."
wait_for_server "$PROXY_PORT"
curl http://127.0.0.1:"$PROXY_PORT"/healthcheck
echo ""
echo "All PD (1P+1D) services are up!"
# Run test in baseline mode
echo "Running PD disagg baseline..."
python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
--service_url "http://localhost:$PROXY_PORT" \
--model_name "$MODEL" \
--mode baseline_pd \
--baseline_file "$BASELINE_PD_FILE" \
"${TEST_FLAGS[@]}"
# Cleanup
echo "Stopping PD (1P+1D) instances..."
for pid in "${PIDS[@]}"; do
kill "$pid" 2>/dev/null || true
done
sleep 2
cleanup_instances
}
# Function to run EPD with 1E + 1P + 1D
run_epd_1e_1p_1d() {
echo "================================"
echo "Running EPD (1E + 1P + 1D)"
echo "================================"
cleanup_instances
rm -rf "$EC_SHARED_STORAGE_PATH"
mkdir -p "$EC_SHARED_STORAGE_PATH"
local ENCODE_PORT=$ENCODE_PORT
local PREFILL_PORT=$PREFILL_PORT
local DECODE_PORT=$DECODE_PORT
local PROXY_PORT=$ENDPOINT_PORT
declare -a PIDS=()
# Start encoder instance
echo "Starting encoder instance on GPU $GPU_E, port $ENCODE_PORT"
env "$DEVICE_AFFINITY_ENV=$GPU_E" vllm serve "$MODEL" \
--port "$ENCODE_PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--enforce-eager \
--gpu-memory-utilization 0.01 \
--enable-request-id-headers \
--no-enable-prefix-caching \
--max-num-batched-tokens 114688 \
--max-num-seqs "$MAX_NUM_SEQS" \
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_producer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}' \
> "$LOG_PATH"/1e1p1d_encoder.log 2>&1 &
PIDS+=($!)
# Start prefill instance
echo "Starting prefill instance on GPU $GPU_P, port $PREFILL_PORT"
env "$DEVICE_AFFINITY_ENV=$GPU_P" \
VLLM_NIXL_SIDE_CHANNEL_PORT=5559 \
vllm serve "$MODEL" \
--port "$PREFILL_PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--enforce-eager \
--enable-mm-embeds \
--gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
--enable-request-id-headers \
--max-num-seqs "$MAX_NUM_SEQS" \
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_consumer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}' \
--kv-transfer-config '{
"kv_connector": "NixlConnector",
"kv_role": "kv_producer"
}' \
> "$LOG_PATH"/1e1p1d_prefill.log 2>&1 &
PIDS+=($!)
# Start decode instance
echo "Starting decode instance on GPU $GPU_D, port $DECODE_PORT"
env "$DEVICE_AFFINITY_ENV=$GPU_D" \
VLLM_NIXL_SIDE_CHANNEL_PORT=6000 \
vllm serve "$MODEL" \
--port "$DECODE_PORT" \
--max-model-len "$MAX_MODEL_LEN" \
--enforce-eager \
--enable-mm-embeds \
--gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
--enable-request-id-headers \
--max-num-seqs "$MAX_NUM_SEQS" \
--allowed-local-media-path "${GIT_ROOT}"/tests/v1/ec_connector/integration \
--kv-transfer-config '{
"kv_connector": "NixlConnector",
"kv_role": "kv_consumer"
}' \
> "$LOG_PATH"/1e1p1d_decode.log 2>&1 &
PIDS+=($!)
# Wait for instances to start
echo "Waiting for encoder instance..."
wait_for_server "$ENCODE_PORT"
echo "Waiting for prefill instance..."
wait_for_server "$PREFILL_PORT"
echo "Waiting for decode instance..."
wait_for_server "$DECODE_PORT"
# Start proxy
echo "Starting EPD proxy on port $PROXY_PORT"
python "${GIT_ROOT}/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" \
--host "0.0.0.0" \
--port "$PROXY_PORT" \
--encode-servers-urls "http://localhost:$ENCODE_PORT" \
--prefill-servers-urls "http://localhost:$PREFILL_PORT" \
--decode-servers-urls "http://localhost:$DECODE_PORT" \
> "$LOG_PATH"/1e1p1d_proxy.log 2>&1 &
PIDS+=($!)
# Wait for proxy
echo "Waiting for proxy..."
wait_for_server "$PROXY_PORT"
curl http://127.0.0.1:"$PROXY_PORT"/v1/models
curl http://127.0.0.1:"$PROXY_PORT"/health
echo ""
echo "All EPD (1E+1P+1D) services are up!"
# Run test in disagg mode
echo "Running EPD (1E+1P+1D) correctness test..."
python "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
--service_url "http://localhost:$PROXY_PORT" \
--model_name "$MODEL" \
--mode disagg \
--baseline_file "$BASELINE_PD_FILE" \
"${TEST_FLAGS[@]}"
# Cleanup
echo "✓✓ 1E+1P+1D Correctness Test finished"
echo "Stopping EPD (1E+1P+1D) instances..."
for pid in "${PIDS[@]}"; do
kill "$pid" 2>/dev/null || true
done
sleep 2
cleanup_instances
}
# Main execution
echo "================================"
echo "EPD Correctness Test Suite"
echo "Model: $MODEL"
echo "================================"
# Step 1: Run baseline
run_baseline
# Step 2: Test 1E + 1PD
run_epd_1e_1pd
# Step 3: Test baseline 1P + 1D
run_baseline_1p_1d
# # Step 4: Test 1E + 1P + 1D
run_epd_1e_1p_1d
# Cleanup output file
rm -f "$BASELINE_FILE"
rm -f "$BASELINE_PD_FILE"
echo "================================"
echo "✓✓ All EPD correctness tests finished!"
echo "================================"
tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
-323
View File
@@ -1,323 +0,0 @@
#!/usr/bin/env bash
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Full-stack EPD validation with ECMooncakeConnector (Mooncake TransferEngine):
# 1) Single-GPU baseline (multimodal) -> saves baseline JSON
# 2) 1 Encoder + 1 PD with Mooncake EC -> compare outputs to baseline
#
# Usage (from repo root):
# ./tests/v1/ec_connector/integration/run_epd_mooncake_ec_full_pipeline.sh
#
# Env:
# MODEL HF model id (default: Qwen/Qwen2.5-VL-3B-Instruct)
# GPU_SINGLE / GPU_E / GPU_PD GPU ids (defaults 0 / 0 / 1)
# ENDPOINT_PORT, ENCODE_PORT, PREFILL_DECODE_PORT
# EC_MOONCAKE_RESERVATION_HOST consumer control host (default 127.0.0.1)
# MOONCAKE_EC_PROTOCOL tcp | rdma (default tcp)
# USE_MM_PROMPTS 1 (default) or 0 for text-only quick sanity
# TIMEOUT_SECONDS wait_for_server timeout (default 1200)
# SKIP_BASELINE set to 1 to reuse existing BASELINE_FILE
# CONCURRENCY / REPEAT concurrent requests / rounds (defaults 3 / 2)
# MAX_MODEL_LEN context length (default 16384)
# E_CUDAGRAPH_MM_ENCODER enable encoder CUDA graphs (default 0)
# E_MM_ENCODER_ATTN_BACKEND optional encoder attention backend override
set -euo pipefail
GIT_ROOT=$(cd "$(dirname "${BASH_SOURCE[0]}")/../../../.." && pwd)
cd "$GIT_ROOT" || exit 1
export PYTHONPATH="${GIT_ROOT}:${PYTHONPATH:-}"
PYTHON_BIN="${PYTHON_BIN:-${GIT_ROOT}/.venv/bin/python}"
MODEL="${MODEL:-Qwen/Qwen2.5-VL-3B-Instruct}"
USE_MM_PROMPTS="${USE_MM_PROMPTS:-1}"
TEST_ARGS=(--concurrency "${CONCURRENCY:-3}" --repeat "${REPEAT:-2}")
if [[ "$USE_MM_PROMPTS" == "1" ]]; then
TEST_ARGS+=(--use_mm_prompts --mm_smoke_test)
fi
MAX_MODEL_LEN="${MAX_MODEL_LEN:-16384}"
E_ATTN_ARGS=()
if [[ -n "${E_MM_ENCODER_ATTN_BACKEND:-}" ]]; then
E_ATTN_ARGS+=(--mm-encoder-attn-backend "$E_MM_ENCODER_ATTN_BACKEND")
fi
case "${E_CUDAGRAPH_MM_ENCODER:-0}" in
0) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":false}' ;;
1) E_COMPILATION_CONFIG='{"cudagraph_mm_encoder":true}' ;;
*) echo "E_CUDAGRAPH_MM_ENCODER must be 0 or 1"; exit 1 ;;
esac
GPU_SINGLE="${GPU_SINGLE:-0}"
GPU_E="${GPU_E:-0}"
GPU_PD="${GPU_PD:-1}"
ENCODE_PORT="${ENCODE_PORT:-19534}"
PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19537}"
ENDPOINT_PORT="${ENDPOINT_PORT:-10002}"
BASELINE_PORT="${BASELINE_PORT:-10003}"
EC_MOONCAKE_RESERVATION_HOST="${EC_MOONCAKE_RESERVATION_HOST:-127.0.0.1}"
EC_MOONCAKE_RESERVATION_PORT="${EC_MOONCAKE_RESERVATION_PORT:-19019}"
MOONCAKE_EC_PROTOCOL="${MOONCAKE_EC_PROTOCOL:-tcp}"
export EC_MOONCAKE_RESERVATION_HOST
export EC_MOONCAKE_RESERVATION_PORT
export MOONCAKE_EC_PROTOCOL
if [[ "$MOONCAKE_EC_PROTOCOL" == "tcp" ]]; then
# TransferEngine may otherwise auto-select RDMA on hosts with an HCA.
export MC_FORCE_TCP=1
else
unset MC_FORCE_TCP
fi
LOG_PATH="${LOG_PATH:-/tmp}"
BASELINE_FILE="${BASELINE_FILE:-/tmp/vllm_epd_mooncake_baseline.txt}"
TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-1200}"
PIDS=()
mkdir -p "$LOG_PATH"
VLLM_SERVE=("$PYTHON_BIN" -m vllm.entrypoints.cli.main serve)
ENC_EC_JSON=$("$PYTHON_BIN" <<PY
import json, os
print(json.dumps({
"ec_connector": "ECMooncakeConnector",
"ec_role": "ec_producer",
"ec_connector_extra_config": {
"mooncake_protocol": os.environ.get("MOONCAKE_EC_PROTOCOL", "tcp"),
},
}, separators=(",", ":")))
PY
)
PD_EC_JSON=$("$PYTHON_BIN" <<PY
import json, os
print(json.dumps({
"ec_connector": "ECMooncakeConnector",
"ec_role": "ec_consumer",
"ec_ip": os.environ["EC_MOONCAKE_RESERVATION_HOST"],
"ec_port": int(os.environ.get("EC_MOONCAKE_RESERVATION_PORT", "19019")),
"ec_connector_extra_config": {
"mooncake_protocol": os.environ.get("MOONCAKE_EC_PROTOCOL", "tcp"),
},
}, separators=(",", ":")))
PY
)
EC_MOONCAKE_RESERVATION_ADDR=$("$PYTHON_BIN" <<'PY'
import os
from vllm.utils.network_utils import make_zmq_path
print(make_zmq_path(
"tcp", os.environ["EC_MOONCAKE_RESERVATION_HOST"],
int(os.environ["EC_MOONCAKE_RESERVATION_PORT"]),
))
PY
)
wait_for_server() {
local port=$1
local pid=$2
local deadline=$((SECONDS + TIMEOUT_SECONDS))
while ((SECONDS < deadline)); do
kill -0 "$pid" 2>/dev/null || return 1
if curl --max-time 2 -fsS "http://localhost:${port}/health" >/dev/null 2>&1; then
return 0
fi
sleep 2
done
return 1
}
cleanup_instances() {
if ((${#PIDS[@]} == 0)); then
return
fi
echo "Cleaning up tracked vLLM / proxy processes..."
for pid in "${PIDS[@]}"; do
kill "$pid" 2>/dev/null || true
done
for _ in {1..10}; do
local alive=0
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
alive=1
fi
done
((alive == 0)) && break
sleep 1
done
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
kill -KILL "$pid" 2>/dev/null || true
fi
wait "$pid" 2>/dev/null || true
done
PIDS=()
}
finish() {
local status=$?
cleanup_instances
if ((status != 0)); then
for log in "${LOG_PATH}"/mooncake_epd_*.log; do
[[ -f "$log" ]] || continue
echo "=== $log (last 100 lines) ==="
tail -100 "$log"
done
fi
exit "$status"
}
trap finish EXIT
trap 'exit 130' INT
trap 'exit 143' TERM
run_baseline() {
echo "================================"
echo "BASELINE (single vLLM, MM if enabled)"
echo "================================"
cleanup_instances
local PORT=$BASELINE_PORT
echo "Starting baseline on GPU $GPU_SINGLE port $PORT"
CUDA_VISIBLE_DEVICES="$GPU_SINGLE" "${VLLM_SERVE[@]}" "$MODEL" \
--port "$PORT" \
--gpu-memory-utilization 0.75 \
--max-num-seqs 32 \
--max-model-len "$MAX_MODEL_LEN" \
--no-enable-prefix-caching \
--limit-mm-per-prompt '{"image":2,"video":0}' \
--allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \
>"${LOG_PATH}/mooncake_epd_baseline.log" 2>&1 &
local BASELINE_PID=$!
PIDS+=("$BASELINE_PID")
echo "Waiting for baseline..."
wait_for_server "$PORT" "$BASELINE_PID" || { echo "Baseline failed to start"; return 1; }
curl -s "http://127.0.0.1:${PORT}/v1/models" | head -c 200 || true
echo ""
"$PYTHON_BIN" "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
--service_url "http://localhost:$PORT" \
--model_name "$MODEL" \
--mode baseline \
--baseline_file "$BASELINE_FILE" \
"${TEST_ARGS[@]}"
cleanup_instances
}
run_epd_mooncake() {
echo "================================"
echo "EPD 1E + 1PD with ECMooncakeConnector"
echo "Reservation port on consumer: $EC_MOONCAKE_RESERVATION_PORT"
echo "Mooncake protocol: $MOONCAKE_EC_PROTOCOL"
echo "================================"
cleanup_instances
echo "Starting ENCODER on GPU $GPU_E port $ENCODE_PORT"
VLLM_SERVER_DEV_MODE=1 CUDA_VISIBLE_DEVICES="$GPU_E" "${VLLM_SERVE[@]}" "$MODEL" \
--port "$ENCODE_PORT" \
--gpu-memory-utilization 0.35 \
--mm-tensor-ipc torch_shm \
--enable-request-id-headers \
--no-enable-prefix-caching \
--max-num-batched-tokens "$MAX_MODEL_LEN" \
--max-num-seqs 32 \
--max-model-len "$MAX_MODEL_LEN" \
--limit-mm-per-prompt '{"image":2,"video":0}' \
--allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \
--ec-transfer-config "$ENC_EC_JSON" \
--compilation-config "$E_COMPILATION_CONFIG" \
"${E_ATTN_ARGS[@]}" \
--worker-extension-cls tests.v1.ec_connector.integration.test_epd_correctness.EncoderGraphTestWorkerExtension \
>"${LOG_PATH}/mooncake_epd_encoder.log" 2>&1 &
local ENCODER_PID=$!
PIDS+=("$ENCODER_PID")
echo "Starting PD on GPU $GPU_PD port $PREFILL_DECODE_PORT"
CUDA_VISIBLE_DEVICES="$GPU_PD" "${VLLM_SERVE[@]}" "$MODEL" \
--port "$PREFILL_DECODE_PORT" \
--gpu-memory-utilization 0.75 \
--mm-tensor-ipc torch_shm \
--enable-mm-embeds \
--enable-request-id-headers \
--max-num-seqs 32 \
--max-model-len "$MAX_MODEL_LEN" \
--no-enable-prefix-caching \
--limit-mm-per-prompt '{"image":2,"video":0}' \
--allowed-local-media-path "${GIT_ROOT}/tests/v1/ec_connector/integration" \
--ec-transfer-config "$PD_EC_JSON" \
>"${LOG_PATH}/mooncake_epd_pd.log" 2>&1 &
local PD_PID=$!
PIDS+=("$PD_PID")
echo "Waiting for encoder..."
wait_for_server "$ENCODE_PORT" "$ENCODER_PID" || { echo "Encoder failed to start"; return 1; }
echo "Waiting for PD..."
wait_for_server "$PREFILL_DECODE_PORT" "$PD_PID" || { echo "PD failed to start"; return 1; }
echo "Starting EPD proxy on $ENDPOINT_PORT"
"$PYTHON_BIN" "${GIT_ROOT}/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py" \
--host "0.0.0.0" \
--port "$ENDPOINT_PORT" \
--encode-servers-urls "http://localhost:$ENCODE_PORT" \
--prefill-servers-urls "disable" \
--decode-servers-urls "http://localhost:$PREFILL_DECODE_PORT" \
--ec-consumer-zmq-addrs \
"$EC_MOONCAKE_RESERVATION_ADDR" \
>"${LOG_PATH}/mooncake_epd_proxy.log" 2>&1 &
local PROXY_PID=$!
PIDS+=("$PROXY_PID")
echo "Waiting for proxy..."
wait_for_server "$ENDPOINT_PORT" "$PROXY_PID" || { echo "Proxy failed to start"; return 1; }
curl -s "http://127.0.0.1:${ENDPOINT_PORT}/health" || true
echo ""
"$PYTHON_BIN" "${GIT_ROOT}/tests/v1/ec_connector/integration/test_epd_correctness.py" \
--service_url "http://localhost:$ENDPOINT_PORT" \
--model_name "$MODEL" \
--mode disagg \
--baseline_file "$BASELINE_FILE" \
"${TEST_ARGS[@]}"
curl -fsS "http://localhost:${ENCODE_PORT}/collective_rpc" \
-H 'Content-Type: application/json' \
-d '{"method":"encoder_graph_stats"}' \
>"${LOG_PATH}/encoder_graph_stats.json"
"$PYTHON_BIN" - "${LOG_PATH}/encoder_graph_stats.json" "${E_CUDAGRAPH_MM_ENCODER:-0}" <<'PY'
import json
import sys
with open(sys.argv[1]) as f:
results = json.load(f)["results"]
assert results, "No encoder workers returned graph statistics"
for stats in results:
if sys.argv[2] == "1":
assert stats["is_captured"], stats
assert stats["graph_hits"] > 0, stats
else:
assert stats["graph_hits"] == 0, stats
print("Encoder graph statistics:", results)
PY
cleanup_instances
}
echo "================================"
echo "1E + 1PD ECMooncake end-to-end correctness"
echo "MODEL=$MODEL"
echo "================================"
if [[ "${SKIP_BASELINE:-0}" != "1" ]]; then
run_baseline
else
echo "SKIP_BASELINE=1 -> using existing $BASELINE_FILE"
[[ -f "$BASELINE_FILE" ]] || { echo "Missing baseline file"; exit 1; }
fi
run_epd_mooncake
echo "================================"
echo "PASS: Mooncake EC EPD matches baseline"
echo "Logs: ${LOG_PATH}/mooncake_epd_*.log"
echo "================================"
tests/v1/ec_connector/integration/test_epd_correctness.py
-395
View File
@@ -1,395 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""EPD Correctness Test.
Tests that EPD (Encoder-Prefill-Decode) disaggregation produces the same
outputs as a baseline single instance.
Usage:
# Baseline mode (saves outputs):
python test_epd_correctness.py \
--service_url http://localhost:8000 \
--model_name Qwen/Qwen2.5-VL-3B-Instruct \
--mode baseline \
--baseline_file .vllm_epd_baseline.txt
# Disagg mode (compares outputs):
python test_epd_correctness.py \
--service_url http://localhost:8000 \
--model_name Qwen/Qwen2.5-VL-3B-Instruct \
--mode disagg \
--baseline_file .vllm_epd_baseline.txt
"""
import argparse
import json
import os
import time
from concurrent.futures import ThreadPoolExecutor
from typing import TYPE_CHECKING
import openai
import requests
from vllm.assets.image import ImageAsset
from vllm.multimodal.utils import encode_image_url
if TYPE_CHECKING:
from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
class EncoderGraphTestWorkerExtension:
"""Expose encoder replay evidence through the existing test RPC endpoint."""
model_runner: "MMEncoderModelRunner"
def encoder_graph_stats(self):
from vllm.v1.worker.mm_encoder_model_runner import MMEncoderModelRunner
runner = self.model_runner
assert isinstance(runner, MMEncoderModelRunner)
manager = runner.model_state.encoder_runner.cudagraph_manager
if manager is None:
return {"is_captured": False, "graph_hits": 0, "graph_misses": 0}
return {"is_captured": manager.is_captured(), **manager.get_cumulative_stats()}
MAX_OUTPUT_LEN = 256
# Sample prompts with multimodal content
image_1 = ImageAsset("stop_sign").pil_image.resize((1280, 720))
image_2 = ImageAsset("cherry_blossom").pil_image.resize((1280, 720))
image_local_path = f"{os.path.dirname(os.path.abspath(__file__))}/hato.jpg"
SAMPLE_PROMPTS_MM: list[dict] = [
{
"messages": [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": encode_image_url(image_1)},
},
{"type": "text", "text": "What's in this image?"},
],
}
],
"description": "Single image query",
},
{
"messages": [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {"url": encode_image_url(image_2)},
},
{
"type": "image_url",
"image_url": {"url": f"file://{image_local_path}"},
},
{"type": "text", "text": "Describe these 2 images in detail."},
],
}
],
"description": "2 images with detailed query",
},
]
# Text-only prompts for mixed testing
SAMPLE_PROMPTS_TEXT: list[dict] = [
{
"messages": [{"role": "user", "content": "What is the capital of France?"}],
"description": "Simple text-only query",
},
{
"messages": [
{"role": "user", "content": "Explain quantum computing in simple terms."}
],
"description": "Text-only explanation request",
},
]
def check_vllm_server(url: str, timeout=5, retries=10) -> bool:
"""Check if the vLLM server is ready.
Args:
url: The URL to check (usually /health or /healthcheck endpoint)
timeout: Timeout in seconds for each request
retries: Number of retries if the server is not ready
Returns:
True if the server is ready, False otherwise
"""
for attempt in range(retries):
try:
response = requests.get(url, timeout=timeout)
if response.status_code == 200:
print(f"Server is ready at {url}")
return True
else:
print(
f"Attempt {attempt + 1}/{retries}: Server returned "
f"status code {response.status_code}"
)
except requests.exceptions.RequestException as e:
print(f"Attempt {attempt + 1}/{retries}: Error connecting: {e}")
time.sleep(2) # Wait before retrying
return False
def run_chat_completion(
base_url: str,
model_name: str,
messages: list,
max_tokens: int = MAX_OUTPUT_LEN,
) -> str:
"""Run a chat completion request.
Args:
base_url: Base URL of the vLLM server
model_name: Name of the model
messages: Messages for chat completion
max_tokens: Maximum tokens to generate
Returns:
Generated text content
"""
with openai.OpenAI(
api_key="EMPTY", base_url=base_url, timeout=120, max_retries=0
) as client:
completion = client.chat.completions.create(
model=model_name,
messages=messages,
max_tokens=max_tokens,
temperature=0.0,
seed=42,
)
content = completion.choices[0].message.content
assert content, "Expected a nonempty completion"
return content
def main():
"""Main test function."""
parser = argparse.ArgumentParser(
description="EPD correctness test - compare disagg vs baseline"
)
parser.add_argument(
"--service_url",
type=str,
required=True,
help="The vLLM service URL (e.g., http://localhost:8000)",
)
parser.add_argument(
"--model_name",
type=str,
required=True,
help="Model name",
)
parser.add_argument(
"--mode",
type=str,
default="baseline",
choices=["baseline", "baseline_pd", "disagg"],
help="Mode: baseline/baseline_pd (saves outputs) or disagg (compares outputs)",
)
parser.add_argument(
"--baseline_file",
type=str,
default=".vllm_epd_baseline.txt",
help="File to save/load baseline outputs",
)
parser.add_argument(
"--use_mm_prompts",
action="store_true",
help="Use multimodal prompts (default: use text-only for quick testing)",
)
parser.add_argument(
"--skip_two_image_prompt",
action="store_true",
help="Skip the two-image multimodal prompt",
)
parser.add_argument(
"--mm_smoke_test",
action="store_true",
help="Use three short, fixed-answer image cases, including duplicate images",
)
parser.add_argument("--concurrency", type=int, default=1)
parser.add_argument("--repeat", type=int, default=1)
args = parser.parse_args()
if args.concurrency < 1 or args.repeat < 1:
parser.error("--concurrency and --repeat must be positive")
if args.mm_smoke_test and (not args.use_mm_prompts or args.skip_two_image_prompt):
parser.error("--mm_smoke_test requires all multimodal prompts")
print(f"Service URL: {args.service_url}")
print(f"Model: {args.model_name}")
print(f"Mode: {args.mode}")
print(f"Output file: {args.baseline_file}")
print(f"Use MM prompts: {args.use_mm_prompts}")
# Determine health check endpoint
if args.mode == "baseline":
health_check_url = f"{args.service_url}/health"
elif args.mode == "baseline_pd":
# Nixl toy proxy use /healthcheck
health_check_url = f"{args.service_url}/healthcheck"
else:
# Disagg EPD proxy uses /health
health_check_url = f"{args.service_url}/health"
if not os.path.exists(args.baseline_file):
raise ValueError(
f"In disagg mode, the output file {args.baseline_file} from "
"baseline does not exist. Run baseline mode first."
)
# Check if server is ready
if not check_vllm_server(health_check_url):
raise RuntimeError(f"vLLM server at {args.service_url} is not ready!")
# Select prompts to use
if args.use_mm_prompts:
test_prompts = (
SAMPLE_PROMPTS_MM[:1] if args.skip_two_image_prompt else SAMPLE_PROMPTS_MM
)
print("Using multimodal prompts")
else:
test_prompts = SAMPLE_PROMPTS_TEXT
print("Using text-only prompts for quick testing")
if args.mm_smoke_test:
image = SAMPLE_PROMPTS_MM[0]["messages"][0]["content"][0]
image_pair = SAMPLE_PROMPTS_MM[1]["messages"][0]["content"][:2]
cases = [
(
"Single image",
[image],
"What word is on the red road sign? Reply with only that word "
"in uppercase.",
"STOP",
),
(
"Two different images",
image_pair,
"Do these pictures contain both flowers and birds? "
"Reply with only YES or NO.",
"YES",
),
(
"Same image twice",
[image, image],
"Read the red road sign in each image, in image order. Reply "
"with only the two uppercase words separated by a comma and a space.",
"STOP, STOP",
),
]
test_prompts = [
{
"description": description,
"expected": expected,
"messages": [
{
"role": "user",
"content": [*images, {"type": "text", "text": question}],
}
],
}
for description, images, question, expected in cases
]
# Run completions
service_url = f"{args.service_url}/v1"
output_strs = {}
def complete(prompt_data):
output = run_chat_completion(
base_url=service_url,
model_name=args.model_name,
messages=prompt_data["messages"],
max_tokens=16 if args.mm_smoke_test else MAX_OUTPUT_LEN,
)
if args.mm_smoke_test:
output = output.strip()
assert output == prompt_data["expected"], (
f"{prompt_data['description']}: expected {prompt_data['expected']!r}, "
f"got {output!r}"
)
return output
# Each round includes concurrent requests sharing image hashes; later rounds
# exercise reuse after previous requests have finished.
with ThreadPoolExecutor(max_workers=args.concurrency) as executor:
for repeat in range(args.repeat):
outputs = executor.map(complete, test_prompts)
for prompt_data, output_str in zip(test_prompts, outputs):
key = prompt_data["description"]
if args.repeat > 1:
key = f"{key} (round {repeat + 1})"
output_strs[key] = output_str
print(f"{key}: {output_str}")
if args.mode in ("baseline", "baseline_pd"):
# Baseline mode: Save outputs
print(f"\nSaving baseline outputs to {args.baseline_file}")
try:
with open(args.baseline_file, "w") as json_file:
json.dump(output_strs, json_file, indent=4)
print("✅ Baseline outputs saved successfully")
except OSError as e:
print(f"Error writing to file: {e}")
raise
else:
# Disagg mode: Load and compare outputs
print(f"\nLoading baseline outputs from {args.baseline_file}")
baseline_outputs = None
try:
with open(args.baseline_file) as json_file:
baseline_outputs = json.load(json_file)
except OSError as e:
print(f"Error reading from file: {e}")
raise
# Verify outputs match
print("\nComparing disagg outputs with baseline...")
assert isinstance(baseline_outputs, dict), "Baseline outputs should be a dict"
assert len(baseline_outputs) == len(output_strs), (
f"Length mismatch: baseline has {len(baseline_outputs)}, "
f"disagg has {len(output_strs)}"
)
all_match = True
for key, baseline_output in baseline_outputs.items():
assert key in output_strs, f"{key} not in disagg outputs"
disagg_output = output_strs[key]
if baseline_output == disagg_output:
print(f"✅ {key}: MATCH")
else:
print(f"❌ {key}: MISMATCH")
print(f" Baseline: {baseline_output}")
print(f" Disagg: {disagg_output}")
all_match = False
assert all_match, "❌❌Disagg outputs do not match baseline!❌❌"
if all_match:
print("\n✅ All outputs match! Test PASSED")
if __name__ == "__main__":
main()
tests/v1/ec_connector/unit/cpu/scheduler/test_embedding_cache.py
-389
View File
@@ -1,389 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for EmbeddingCache.
Pure integer bookkeeping — no mmap, torch, or CUDA — so these run anywhere.
"""
import threading
import pytest
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
EmbeddingCache,
)
def _cache(num_blocks: int = 8) -> EmbeddingCache:
return EmbeddingCache(num_blocks)
# ── alloc ────────────────────────────────────────────────────────────────────
def test_alloc_returns_unique_ascending_block_ids():
cache = _cache()
entry = cache.alloc("a", 3)
assert entry is not None
assert len(entry.block_ids) == 3
assert len(set(entry.block_ids)) == 3
assert all(0 <= i < 8 for i in entry.block_ids)
# Ascending order lets consumers detect which blocks are adjacent.
assert list(entry.block_ids) == sorted(entry.block_ids)
def test_alloc_returns_tuple():
cache = _cache()
entry = cache.alloc("a", 2)
assert isinstance(entry.block_ids, tuple)
def test_alloc_returns_none_when_full():
cache = _cache(num_blocks=4)
cache.alloc("a", 4)
assert cache.alloc("b", 1) is None
def test_alloc_asserts_duplicate_key():
cache = _cache()
cache.alloc("a", 2)
with pytest.raises(AssertionError, match="duplicate alloc"):
cache.alloc("a", 2)
def test_alloc_asserts_n_blocks_exceeds_capacity():
cache = _cache(num_blocks=4)
with pytest.raises(AssertionError, match="capacity"):
cache.alloc("a", 5)
def test_alloc_zero_blocks():
cache = _cache()
entry = cache.alloc("a", 0)
assert entry is not None
assert entry.block_ids == ()
def test_alloc_starts_not_ready():
cache = _cache()
entry = cache.alloc("a", 2)
assert not entry.ready
assert not entry.evictable
# ── get ──────────────────────────────────────────────────────────────────────
def test_get_missing_returns_none():
cache = _cache()
assert cache.get("a") is None
def test_get_returns_entry():
cache = _cache()
alloc_entry = cache.alloc("a", 2)
got = cache.get("a")
assert got is alloc_entry
def test_get_returns_not_ready_entry():
cache = _cache()
cache.alloc("a", 2)
entry = cache.get("a")
assert entry is not None
assert not entry.ready
# ── mark_ready ───────────────────────────────────────────────────────────────
def test_mark_ready_transitions():
cache = _cache()
cache.alloc("a", 2)
entry = cache.get("a")
assert not entry.ready
cache.mark_ready("a")
assert entry.ready
assert entry.evictable
def test_mark_ready_asserts_already_ready():
cache = _cache()
cache.alloc("a", 2)
cache.mark_ready("a")
with pytest.raises(AssertionError, match="already-ready"):
cache.mark_ready("a")
def test_mark_ready_asserts_missing_key():
cache = _cache()
with pytest.raises(KeyError):
cache.mark_ready("nonexistent")
# ── pin / unpin ──────────────────────────────────────────────────────────────
def test_pin_asserts_not_ready():
cache = _cache()
cache.alloc("a", 2)
with pytest.raises(AssertionError, match="not-ready"):
cache.pin("a")
def test_pin_unpin_basic():
cache = _cache()
cache.alloc("a", 2)
cache.mark_ready("a")
cache.pin("a")
entry = cache.get("a")
assert not entry.evictable
cache.unpin("a")
assert entry.evictable
def test_nested_pin():
cache = _cache()
cache.alloc("a", 2)
cache.mark_ready("a")
cache.pin("a")
cache.pin("a")
cache.unpin("a")
entry = cache.get("a")
assert not entry.evictable # still pinned (count=1)
cache.unpin("a")
assert entry.evictable
def test_unpin_asserts_on_unpinned():
cache = _cache()
cache.alloc("a", 2)
cache.mark_ready("a")
with pytest.raises(AssertionError, match="unpinned"):
cache.unpin("a")
# ── has_held_entries ─────────────────────────────────────────────────────────
def test_has_held_entries_empty_cache():
cache = _cache()
assert cache.has_held_entries() is False
def test_has_held_entries_true_while_not_ready():
cache = _cache()
cache.alloc("a", 2) # not ready → held
assert cache.has_held_entries() is True
cache.mark_ready("a") # ready + unpinned → settled
assert cache.has_held_entries() is False
def test_has_held_entries_true_while_pinned():
cache = _cache()
cache.alloc("a", 2)
cache.mark_ready("a")
assert cache.has_held_entries() is False
cache.pin("a") # pinned → held
assert cache.has_held_entries() is True
cache.unpin("a")
assert cache.has_held_entries() is False
# ── eviction ─────────────────────────────────────────────────────────────────
def test_eviction_frees_ready_unpinned_fifo():
cache = _cache(num_blocks=4)
cache.alloc("a", 2)
cache.mark_ready("a")
cache.alloc("b", 2)
cache.mark_ready("b")
# Pool full. Alloc should evict "a" (oldest).
entry_c = cache.alloc("c", 2)
assert entry_c is not None
assert cache.get("a") is None # evicted
assert cache.get("b") is not None # still there
def test_eviction_skips_pinned():
cache = _cache(num_blocks=4)
cache.alloc("a", 2)
cache.mark_ready("a")
cache.pin("a")
cache.alloc("b", 2)
cache.mark_ready("b")
# "a" is pinned, so "b" should be evicted instead.
entry_c = cache.alloc("c", 2)
assert entry_c is not None
assert cache.get("a") is not None # pinned, survived
assert cache.get("b") is None # evicted
cache.unpin("a")
def test_eviction_skips_not_ready():
cache = _cache(num_blocks=4)
cache.alloc("a", 2) # not ready
cache.alloc("b", 2)
cache.mark_ready("b")
# "a" is not ready (non-evictable), "b" is evictable.
entry_c = cache.alloc("c", 2)
assert entry_c is not None
assert cache.get("a") is not None # not ready, survived
assert cache.get("b") is None # evicted
cache.mark_ready("a") # clean up
def test_eviction_returns_none_if_all_non_evictable():
cache = _cache(num_blocks=4)
cache.alloc("a", 2)
cache.mark_ready("a")
cache.pin("a")
cache.alloc("b", 2) # not ready
# Both entries are non-evictable (a=pinned, b=not-ready).
assert cache.alloc("c", 1) is None
# Fail-fast: nothing was evicted.
assert cache.get("a") is not None
assert cache.get("b") is not None
cache.unpin("a")
def test_fail_fast_no_wasted_eviction():
cache = _cache(num_blocks=6)
cache.alloc("a", 2)
cache.mark_ready("a") # evictable: 2 blocks
cache.alloc("b", 2) # not ready
cache.alloc("c", 2) # not ready
# free=0, evictable=2, need 3 -> cannot succeed
assert cache.alloc("d", 3) is None
# "a" was NOT evicted despite being evictable
assert cache.get("a") is not None
def test_eviction_multiple_entries():
cache = _cache(num_blocks=6)
cache.alloc("a", 2)
cache.mark_ready("a")
cache.alloc("b", 2)
cache.mark_ready("b")
cache.alloc("c", 2)
cache.mark_ready("c")
# Need 4 blocks: must evict both "a" and "b".
entry_d = cache.alloc("d", 4)
assert entry_d is not None
assert cache.get("a") is None
assert cache.get("b") is None
assert cache.get("c") is not None
# ── thread safety ────────────────────────────────────────────────────────────
def test_concurrent_pin_unpin():
cache = _cache(num_blocks=8)
cache.alloc("a", 4)
cache.mark_ready("a")
errors: list[Exception] = []
barrier = threading.Barrier(8)
def worker():
barrier.wait()
try:
for _ in range(500):
cache.pin("a")
cache.unpin("a")
except Exception as exc:
errors.append(exc)
threads = [threading.Thread(target=worker) for _ in range(8)]
for t in threads:
t.start()
for t in threads:
t.join()
assert not errors, errors
entry = cache.get("a")
assert entry.evictable # all pins balanced
def test_concurrent_alloc_evict():
cache = _cache(num_blocks=4)
errors: list[Exception] = []
barrier = threading.Barrier(4)
counter = {"i": 0}
counter_lock = threading.Lock()
def worker():
barrier.wait()
try:
for _ in range(200):
with counter_lock:
key = f"k{counter['i']}"
counter["i"] += 1
entry = cache.alloc(key, 1)
if entry is not None:
cache.mark_ready(key)
except Exception as exc:
errors.append(exc)
threads = [threading.Thread(target=worker) for _ in range(4)]
for t in threads:
t.start()
for t in threads:
t.join()
assert not errors, errors
# ── discard ──────────────────────────────────────────────────────────────────
def test_discard_not_ready_returns_blocks():
cache = EmbeddingCache(num_blocks=4)
entry = cache.alloc("h", 3)
assert entry is not None
cache.discard("h")
assert cache.get("h") is None
# All blocks reclaimed: a fresh 4-block alloc now fits.
assert cache.alloc("h2", 4) is not None
def test_discard_rejects_ready_entry():
cache = EmbeddingCache(num_blocks=4)
cache.alloc("h", 1)
cache.mark_ready("h")
with pytest.raises(AssertionError):
cache.discard("h")
def test_discard_rejects_absent_key():
cache = EmbeddingCache(num_blocks=4)
with pytest.raises(KeyError):
cache.discard("nope")
# ── pin_if_ready ─────────────────────────────────────────────────────────────
def test_pin_if_ready_pins_and_returns_entry():
cache = EmbeddingCache(num_blocks=4)
cache.alloc("h", 2)
cache.mark_ready("h")
entry = cache.pin_if_ready("h")
assert entry is not None and entry.ready and len(entry.block_ids) == 2
assert cache.alloc("h2", 4) is None # pinned ⇒ not evictable
def test_pin_if_ready_none_when_absent():
cache = EmbeddingCache(num_blocks=4)
assert cache.pin_if_ready("nope") is None
def test_pin_if_ready_returns_unpinned_entry_when_not_ready():
"""The producer needs "present but not ready" told apart from "absent":
the first is NACKed as retryable, the second as a miss."""
cache = EmbeddingCache(num_blocks=4)
cache.alloc("h", 1) # not-ready
entry = cache.pin_if_ready("h")
assert entry is not None and not entry.ready
assert cache.alloc("h2", 4) is None # not-ready ⇒ still not evictable
tests/v1/ec_connector/unit/cpu/scheduler/test_scheduler.py
-466
View File
@@ -1,466 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import uuid
import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
from vllm.config.ec_transfer import ECRole
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUWorkerMetadata,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
ECSharedRegion,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
_N = 16
_BS = 64
class _Pos:
def __init__(self, offset, length):
self.offset = offset
self.length = length
def get_num_embeds(self):
return self.length
class _Feature:
def __init__(self, mm_hash, length=1, identifier=None):
self.mm_hash = mm_hash
self.identifier = identifier if identifier is not None else mm_hash
self.mm_position = _Pos(0, length)
self.modality = "image"
class _Request:
_counter = 0
def __init__(self, features, request_id=None):
self.mm_features = features
if request_id is None:
_Request._counter += 1
self.request_id = f"req_{_Request._counter}"
else:
self.request_id = request_id
class _WorkerOutput:
"""Stand-in for ECConnectorOutput carrying a completion report.
`saves` are mm_hashes; `loads` are transfer ids. One entry per reporting
rank, so a list holding the same id twice is two ranks reporting it.
"""
def __init__(self, *, saves=None, loads=None):
self.ec_connector_worker_meta = ECCPUWorkerMetadata(
completed_saves=list(saves or []),
completed_loads=list(loads or []),
)
def _make_scheduler(
monkeypatch,
num_blocks=_N,
*,
ec_role: ECRole = "ec_both",
tensor_parallel_size=1,
) -> ECCPUScheduler:
region = ECSharedRegion(
engine_id=str(uuid.uuid4()),
num_blocks=num_blocks,
block_size_bytes=_BS,
)
monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)
return ECCPUScheduler(
create_ec_vllm_config(
ec_role=ec_role,
tensor_parallel_size=tensor_parallel_size,
)
)
def _load_ids(meta) -> list[int]:
"""Transfer ids the scheduler dispatched in this step's metadata."""
return [transfer_id for transfer_id, _, _ in meta.loads.values()]
def _load_blocks(meta, mm_hash: str) -> list[int]:
"""Block ids the scheduler dispatched for one load."""
return meta.loads[mm_hash][1]
def _seed_cached(s: ECCPUScheduler, mm_hash: str, n_blocks: int):
"""Pre-populate a ready cache entry backed by real blocks."""
s._cache.alloc(mm_hash, n_blocks, (n_blocks, _BS // s._dtype.itemsize))
s._cache.mark_ready(mm_hash)
def test_offload_reuse_cycle(monkeypatch):
s = _make_scheduler(monkeypatch)
# Two blocks, so the reload also pins down that every allocated block id
# comes back in meta.loads, not just the first.
req = _Request([_Feature("h1", length=2)])
# Step A: first sight — allocate and save.
assert s.has_cache_item("h1") is False
assert s.ensure_cache_available(req, 0) is True
s.update_state_after_alloc(req, 0)
meta_a = s.build_connector_meta(scheduler_output=None)
assert "h1" in meta_a.saves
assert meta_a.loads == {}
# Not ready until the worker reports the save memcpy complete.
assert s.has_cache_item("h1") is False
# Worker reports the save copy finished → entry becomes ready.
s.update_connector_output(_WorkerOutput(saves=["h1"]))
assert s.has_cache_item("h1") is True
# Step C: reload now works — entry is ready at update_state_after_alloc time.
s.update_state_after_alloc(req, 0)
meta_c = s.build_connector_meta(scheduler_output=None)
assert "h1" in meta_c.loads
assert _load_blocks(meta_c, "h1") == meta_a.saves["h1"]
s.shutdown()
def test_has_cache_item_false_when_not_consumer(monkeypatch):
s = _make_scheduler(monkeypatch, ec_role="ec_producer")
assert s.has_cache_item("anything") is False
s.shutdown()
def test_connector_keys_on_identifier_not_mm_hash(monkeypatch):
"""The connector must key the encoder cache on feature.identifier (what
has_cache_item is called with), NOT feature.mm_hash."""
s = _make_scheduler(monkeypatch)
req = _Request([_Feature("PROC_KEY", length=1, identifier="ENC_KEY")])
s.update_state_after_alloc(req, 0)
meta = s.build_connector_meta(scheduler_output=None)
assert "ENC_KEY" in meta.saves
assert "PROC_KEY" not in meta.saves
# Not ready until the completion report.
assert s.has_cache_item("ENC_KEY") is False
# After the report it becomes ready — keyed on identifier.
s.update_connector_output(_WorkerOutput(saves=["ENC_KEY"]))
assert s.has_cache_item("ENC_KEY") is True
assert s.has_cache_item("PROC_KEY") is False
s.shutdown()
def test_cpu_region_fifo_eviction(monkeypatch):
# Region holds exactly 2 one-block encodings.
s = _make_scheduler(monkeypatch, num_blocks=2)
# Fill with A and B, marking each ready via a completion report.
for h in ("A", "B"):
s.update_state_after_alloc(_Request([_Feature(h, length=1)]), 0)
s.build_connector_meta(scheduler_output=None)
s.update_connector_output(_WorkerOutput(saves=[h]))
assert s.has_cache_item("A") is True
assert s.has_cache_item("B") is True
# Allocate C: must evict A (oldest) to make room.
s.update_state_after_alloc(_Request([_Feature("C", length=1)]), 0)
s.build_connector_meta(scheduler_output=None)
assert s.has_cache_item("A") is False # evicted
assert s.has_cache_item("B") is True
def test_multiple_mm_items_per_request(monkeypatch):
s = _make_scheduler(monkeypatch)
req = _Request([_Feature("h1", length=1), _Feature("h2", length=1)])
s.update_state_after_alloc(req, 0)
s.update_state_after_alloc(req, 1)
meta = s.build_connector_meta(scheduler_output=None)
assert set(meta.saves) == {"h1", "h2"}
# Worker reports both saves complete.
s.update_connector_output(_WorkerOutput(saves=["h1", "h2"]))
assert s.has_cache_item("h1") is True
assert s.has_cache_item("h2") is True
# Now reload both.
s.update_state_after_alloc(req, 0)
s.update_state_after_alloc(req, 1)
meta2 = s.build_connector_meta(scheduler_output=None)
assert set(meta2.loads) == {"h1", "h2"}
s.shutdown()
def test_loads_only_serves_hashes_touched_this_step(monkeypatch):
s = _make_scheduler(monkeypatch)
_seed_cached(s, "a", n_blocks=2)
_seed_cached(s, "b", n_blocks=2)
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
meta = s.build_connector_meta(scheduler_output=None)
assert "a" in meta.loads
assert "b" not in meta.loads
s.shutdown()
def test_repeated_reload_same_step_loads_once(monkeypatch):
"""The same mm_hash requested twice in one step must appear in
meta.loads exactly once."""
s = _make_scheduler(monkeypatch)
_seed_cached(s, "a", n_blocks=2)
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
meta = s.build_connector_meta(scheduler_output=None)
assert list(meta.loads.keys()).count("a") == 1
s.shutdown()
def test_load_not_emitted_for_uncached_entry(monkeypatch):
"""A feature that was never saved must not appear in meta.loads."""
s = _make_scheduler(monkeypatch)
s.update_state_after_alloc(_Request([_Feature("missing")]), 0)
meta = s.build_connector_meta(scheduler_output=None)
assert meta.loads == {}
s.shutdown()
def test_load_pin_protects_blocks_until_completion(monkeypatch):
"""Loaded blocks stay pinned until the worker reports the load memcpy
complete. Observable: a save that needs eviction cannot reclaim a
still-pinned loaded entry, but can once the completion report arrives."""
# 2 blocks, both occupied by "a" (ready).
s = _make_scheduler(monkeypatch, num_blocks=2)
_seed_cached(s, "a", n_blocks=2)
# Load "a" — pins it.
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
meta = s.build_connector_meta(scheduler_output=None)
assert "a" in meta.loads
(transfer_id,) = _load_ids(meta)
# Try to save "b" (needs 2 blocks, but "a" is pinned) — must fail.
s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
meta2 = s.build_connector_meta(scheduler_output=None)
assert "b" not in meta2.saves # can't evict pinned "a"
# The one participant reports the load complete → "a" unpins. Now "b" can
# evict it.
s.update_connector_output(_WorkerOutput(loads=[transfer_id]))
s.update_state_after_alloc(_Request([_Feature("b", length=2)]), 0)
meta3 = s.build_connector_meta(scheduler_output=None)
assert "b" in meta3.saves
s.shutdown()
def _can_evict(s: ECCPUScheduler, mm_hash: str, n_blocks: int) -> bool:
"""Whether a new save of `mm_hash` can claim `n_blocks`, i.e. whether the
blocks currently held by other entries are reclaimable."""
s.update_state_after_alloc(_Request([_Feature(mm_hash, length=n_blocks)]), 0)
meta = s.build_connector_meta(scheduler_output=None)
return mm_hash in meta.saves
def test_load_needs_every_participant_before_unpin(monkeypatch):
"""With two participating ranks, one report must not release the pin.
This is the multi-rank case: each rank copies the same CPU blocks into its
own GPU memory, so the blocks stay in use until the last rank is done.
"""
s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2)
_seed_cached(s, "a", n_blocks=2)
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
(transfer_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))
# First rank reports: still in use by the second, so "a" must hold.
s.update_connector_output(_WorkerOutput(loads=[transfer_id]))
assert _can_evict(s, "b", 2) is False
assert s.has_pending_push_work() is True
# Second rank reports: now it is genuinely free.
s.update_connector_output(_WorkerOutput(loads=[transfer_id]))
assert s.has_pending_push_work() is False
assert _can_evict(s, "c", 2) is True
s.shutdown()
def test_late_report_does_not_release_a_later_load_of_same_hash(monkeypatch):
"""A straggling report from one dispatch must not release the pin taken by
a subsequent dispatch of the same mm_hash.
This is what the transfer id is for: mm_hash identifies a cache entry, not
a transfer, so it cannot tell the two dispatches apart.
"""
s = _make_scheduler(monkeypatch, num_blocks=2, tensor_parallel_size=2)
_seed_cached(s, "a", n_blocks=2)
# First dispatch, fully reported by both ranks → pin released.
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
(first_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))
s.update_connector_output(_WorkerOutput(loads=[first_id, first_id]))
assert s.has_pending_push_work() is False
# Second dispatch of the same hash takes a fresh pin.
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
(second_id,) = _load_ids(s.build_connector_meta(scheduler_output=None))
assert second_id != first_id
# A replayed report for the first dispatch must be ignored entirely.
s.update_connector_output(_WorkerOutput(loads=[first_id, first_id]))
assert s.has_pending_push_work() is True
assert _can_evict(s, "b", 2) is False
s.shutdown()
def test_region_full_skips_save_and_never_blocks(monkeypatch):
"""When the region is fully occupied by pinned entries, new saves are
silently skipped and ensure_cache_available never blocks."""
s = _make_scheduler(monkeypatch, num_blocks=1)
_seed_cached(s, "pinned", n_blocks=1)
s._cache.pin("pinned")
req = _Request([_Feature("new", length=1)])
assert s.ensure_cache_available(req, 0) is True
s.update_state_after_alloc(req, 0)
meta = s.build_connector_meta(scheduler_output=None)
assert "new" not in meta.saves
assert s.has_cache_item("new") is False
s._cache.unpin("pinned")
s.shutdown()
def test_producer_only_never_emits_loads(monkeypatch):
"""A producer-only scheduler must never populate meta.loads, even when
entries are ready."""
s = _make_scheduler(monkeypatch, ec_role="ec_producer")
req = _Request([_Feature("h1", length=1)])
s.update_state_after_alloc(req, 0)
meta = s.build_connector_meta(scheduler_output=None)
assert "h1" in meta.saves
assert meta.loads == {}
s.update_connector_output(_WorkerOutput(saves=["h1"]))
# Even if we re-encounter the same feature, no load.
s.update_state_after_alloc(req, 0)
meta2 = s.build_connector_meta(scheduler_output=None)
assert meta2.loads == {}
s.shutdown()
def test_consumer_only_never_emits_saves(monkeypatch):
"""A consumer-only scheduler must never populate meta.saves."""
s = _make_scheduler(monkeypatch, ec_role="ec_consumer")
_seed_cached(s, "a", n_blocks=1)
req = _Request([_Feature("a", length=1)])
s.update_state_after_alloc(req, 0)
meta = s.build_connector_meta(scheduler_output=None)
assert meta.saves == {}
assert "a" in meta.loads
s.shutdown()
def test_consumer_only_has_cache_item(monkeypatch):
"""Consumer-only: has_cache_item works normally for ready entries."""
s = _make_scheduler(monkeypatch, ec_role="ec_consumer")
assert s.has_cache_item("a") is False
_seed_cached(s, "a", n_blocks=1)
assert s.has_cache_item("a") is True
s.shutdown()
def test_save_not_emitted_for_already_cached_entry(monkeypatch):
"""An entry already in the cache (from a prior save) must not trigger
a second allocation or appear in meta.saves again."""
s = _make_scheduler(monkeypatch)
req = _Request([_Feature("h1", length=1)])
# First encounter → save.
s.update_state_after_alloc(req, 0)
meta1 = s.build_connector_meta(scheduler_output=None)
assert "h1" in meta1.saves
# Second encounter (same step or later) → no save.
s.update_state_after_alloc(req, 0)
meta2 = s.build_connector_meta(scheduler_output=None)
assert "h1" not in meta2.saves
s.shutdown()
# ── completion reporting (update_connector_output) ──────────────────────────
def test_update_connector_output_ignores_foreign_meta(monkeypatch):
"""A worker output whose payload is not an ECCPUWorkerMetadata is a
no-op — no cache mutation, no crash."""
s = _make_scheduler(monkeypatch)
_seed_cached(s, "a", n_blocks=1)
class _Foreign:
ec_connector_worker_meta = object()
s.update_connector_output(_Foreign()) # must not raise
assert s.has_cache_item("a") is True
s.shutdown()
def test_update_connector_output_ignores_unknown_report(monkeypatch):
"""Reports for a never-seen save hash or a transfer id this scheduler never
dispatched are dropped rather than mutating unrelated state."""
s = _make_scheduler(monkeypatch)
s.update_connector_output(_WorkerOutput(saves=["gone"], loads=[4242]))
assert s.has_cache_item("gone") is False
s.shutdown()
# ── has_pending_push_work ───────────────────────────────────────────────────
def test_has_pending_push_work_tracks_inflight_save(monkeypatch):
"""A dispatched-but-unconfirmed save keeps push work pending until the
completion report marks it ready."""
s = _make_scheduler(monkeypatch)
assert s.has_pending_push_work() is False
s.update_state_after_alloc(_Request([_Feature("h1", length=1)]), 0)
s.build_connector_meta(scheduler_output=None)
assert s.has_pending_push_work() is True # save not yet confirmed
s.update_connector_output(_WorkerOutput(saves=["h1"]))
assert s.has_pending_push_work() is False
s.shutdown()
def test_has_pending_push_work_tracks_inflight_load(monkeypatch):
"""A dispatched-but-unconfirmed load keeps push work pending until the
unpin completion report arrives."""
s = _make_scheduler(monkeypatch)
_seed_cached(s, "a", n_blocks=1)
assert s.has_pending_push_work() is False # ready + unpinned
s.update_state_after_alloc(_Request([_Feature("a")]), 0)
meta = s.build_connector_meta(scheduler_output=None)
assert s.has_pending_push_work() is True # pinned, awaiting unpin
s.update_connector_output(_WorkerOutput(loads=_load_ids(meta)))
assert s.has_pending_push_work() is False
s.shutdown()
# ── shutdown ────────────────────────────────────────────────────────────────
def test_shutdown_disables_roles_and_cleans_region(monkeypatch):
"""After shutdown the scheduler serves no items and the region is
cleaned up."""
s = _make_scheduler(monkeypatch, num_blocks=2)
_seed_cached(s, "a", n_blocks=2)
assert s.has_cache_item("a") is True
s.shutdown()
assert s.has_cache_item("a") is False # consumer role disabled
tests/v1/ec_connector/unit/cpu/test_connector.py
-71
View File
@@ -1,71 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from unittest.mock import MagicMock
import pytest
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
from vllm.distributed.ec_transfer.ec_connector.cpu.connector import ECCPUConnector
from vllm.distributed.ec_transfer.ec_connector.factory import ECConnectorFactory
def _cfg():
ec = MagicMock()
ec.is_ec_producer = True
ec.is_ec_consumer = True
cfg = MagicMock()
cfg.ec_transfer_config = ec
return cfg
def test_scheduler_role_builds_only_scheduler(monkeypatch):
fake_sched = MagicMock()
monkeypatch.setattr(ECCPUConnector, "_make_scheduler", lambda self, cfg: fake_sched)
c = ECCPUConnector(_cfg(), ECConnectorRole.SCHEDULER)
assert c.connector_scheduler is fake_sched
assert c.connector_worker is None
c.has_cache_item("x")
fake_sched.has_cache_item.assert_called_once_with("x")
def test_worker_role_builds_only_worker(monkeypatch):
fake_worker = MagicMock()
monkeypatch.setattr(ECCPUConnector, "_make_worker", lambda self, cfg: fake_worker)
c = ECCPUConnector(_cfg(), ECConnectorRole.WORKER)
assert c.connector_worker is fake_worker
assert c.connector_scheduler is None
@pytest.mark.parametrize("role", [ECConnectorRole.SCHEDULER, ECConnectorRole.WORKER])
def test_v1_model_runner_is_rejected(monkeypatch, role):
"""The V1 runner never reports save completions, so refuse to start."""
monkeypatch.setattr(
ECCPUConnector, "_make_scheduler", lambda self, cfg: MagicMock()
)
monkeypatch.setattr(ECCPUConnector, "_make_worker", lambda self, cfg: MagicMock())
cfg = _cfg()
cfg.use_v2_model_runner = False
with pytest.raises(ValueError, match="V2 model runner"):
ECCPUConnector(cfg, role)
def test_factory_registered():
cls = ECConnectorFactory._registry["ECCPUConnector"]()
assert cls is ECCPUConnector
def test_request_finished_forwards_to_scheduler(monkeypatch):
from unittest.mock import MagicMock
import vllm.distributed.ec_transfer.ec_connector.cpu.connector as conn_mod # noqa: F401
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
from vllm.distributed.ec_transfer.ec_connector.cpu.connector import ECCPUConnector
fake_sched = MagicMock()
fake_sched.request_finished.return_value = (False, {"h1": {"peer_port": 1}})
monkeypatch.setattr(ECCPUConnector, "_make_scheduler", lambda self, cfg: fake_sched)
c = ECCPUConnector(_cfg(), ECConnectorRole.SCHEDULER)
req = MagicMock()
assert c.request_finished(req) == (False, {"h1": {"peer_port": 1}})
fake_sched.request_finished.assert_called_once_with(req)
tests/v1/ec_connector/unit/cpu/test_ec_shared_region.py
-396
View File
@@ -1,396 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for ECSharedRegion (the mmap substrate)."""
import contextlib
import ctypes
import errno
import mmap
import os
import uuid
from unittest.mock import MagicMock, patch
import pytest
import torch
from vllm.config import ParallelConfig
from vllm.distributed.ec_transfer.ec_connector.cpu import common
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
ECSharedRegion,
_wait_for_file_size,
)
def _make_region(num_blocks: int = 8, block_size_bytes: int = 64) -> ECSharedRegion:
return ECSharedRegion(
engine_id=str(uuid.uuid4()),
num_blocks=num_blocks,
block_size_bytes=block_size_bytes,
)
@pytest.fixture
def region() -> ECSharedRegion:
r = _make_region()
yield r
r.cleanup()
@contextlib.contextmanager
def _region(**kwargs):
"""Context manager: create one region, clean up on exit."""
r = _make_region(**kwargs)
try:
yield r
finally:
r.cleanup()
def _page_residency(mmap_obj: mmap.mmap, length: int) -> list[bool]:
"""Return Linux page-residency bits for a writable mmap."""
try:
libc = ctypes.CDLL(None, use_errno=True)
mincore = libc.mincore
except AttributeError:
pytest.skip("mincore is unavailable")
page_count = (length + mmap.PAGESIZE - 1) // mmap.PAGESIZE
mincore.argtypes = [
ctypes.c_void_p,
ctypes.c_size_t,
ctypes.POINTER(ctypes.c_ubyte),
]
mincore.restype = ctypes.c_int
vector = (ctypes.c_ubyte * page_count)()
address = ctypes.addressof(ctypes.c_ubyte.from_buffer(mmap_obj))
result = mincore(ctypes.c_void_p(address), length, vector)
if result != 0:
raise OSError(ctypes.get_errno())
return [bool(value & 1) for value in vector]
# ── mmap sharing between two instances ───────────────────────────────────────
def test_second_instance_opens_existing_file_and_shares_memory():
instance_id = str(uuid.uuid4())
r1 = ECSharedRegion(
engine_id=instance_id,
num_blocks=4,
block_size_bytes=64,
)
try:
r2 = ECSharedRegion(
engine_id=instance_id,
num_blocks=4,
block_size_bytes=64,
)
try:
assert r1._is_creator
assert not r2._is_creator
# Both map the same physical pages; writes via r1 are visible via r2.
r1.blocks[0, :4] = torch.tensor([10, 20, 30, 40], dtype=torch.int8)
assert r2.blocks[0, :4].tolist() == [10, 20, 30, 40]
finally:
r2.cleanup()
finally:
r1.cleanup()
def test_dense_dp_engines_get_distinct_regions():
"""Dense DP engines reset data_parallel_rank to 0 but share instance_id,
so the region must be keyed by data_parallel_index."""
instance_id = str(uuid.uuid4())
regions = []
try:
for dp_index in range(2):
parallel_config = ParallelConfig(
data_parallel_size=2, data_parallel_rank=dp_index
)
parallel_config.reconfigure_for_independent_dp_rank()
cfg = MagicMock()
cfg.instance_id = instance_id
cfg.parallel_config = parallel_config
cfg.model_config.dtype = torch.float16
cfg.model_config.get_inputs_embeds_size.return_value = 8
cfg.ec_transfer_config.ec_connector_extra_config = {"ec_cpu_bytes": 64}
regions.append(common.create_ec_shared_region(cfg))
assert all(r._is_creator for r in regions)
assert regions[0]._mmap_path != regions[1]._mmap_path
finally:
for r in regions:
r.cleanup()
def test_only_creator_unlinks_file_on_cleanup():
"""Critical contract: if the non-creator unlinks, the creator's mmap path
becomes a dangling backing file and a third opener would create a new one
out from under the creator."""
instance_id = str(uuid.uuid4())
r1 = ECSharedRegion(
engine_id=instance_id,
num_blocks=4,
block_size_bytes=64,
)
path = r1._mmap_path
r2 = ECSharedRegion(
engine_id=instance_id,
num_blocks=4,
block_size_bytes=64,
)
# Non-creator goes away first — file must still be on disk for r1.
r2.cleanup()
assert os.path.exists(path), "non-creator cleanup must not unlink the file"
# Creator goes away — file is removed.
r1.cleanup()
assert not os.path.exists(path), "creator cleanup must unlink the file"
# ── _wait_for_file_size ──────────────────────────────────────────────────────
def test_wait_for_file_size_returns_when_already_big_enough(tmp_path):
"""The fast path: file already at expected size — return immediately."""
p = tmp_path / "f.bin"
p.write_bytes(b"\x00" * 128)
fd = os.open(str(p), os.O_RDONLY)
try:
_wait_for_file_size(fd, expected_size=128, timeout=1.0) # must not raise
finally:
os.close(fd)
def test_wait_for_file_size_times_out_when_file_stays_empty(tmp_path):
p = tmp_path / "f.bin"
p.write_bytes(b"")
fd = os.open(str(p), os.O_RDONLY)
try:
with pytest.raises(TimeoutError):
_wait_for_file_size(fd, expected_size=4096, timeout=0.05)
finally:
os.close(fd)
# ── MADV_POPULATE_WRITE pre-faulting / fallback ──────────────────────────────
def test_madvise_success_selects_madvise_population(monkeypatch):
"""A successful probe must keep using MADV_POPULATE_WRITE — the fallback
helper must never be invoked on a kernel that accepts the advice.
Spies on both helpers so we can verify call counts: 1 probe call over a
single page + 1 populate call over the whole flat region, fallback 0 times.
"""
from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
madvise_calls: list[tuple[int, int, int]] = []
fallback_calls: list[tuple[int, int]] = []
def _spy_madvise(mm, off, ln):
madvise_calls.append((off, ln, id(mm)))
def _spy_fallback(mm, off, ln):
fallback_calls.append((off, ln))
monkeypatch.setattr(esr, "_madvise_populate_write", _spy_madvise)
monkeypatch.setattr(esr, "_fallback_populate_write", _spy_fallback)
num_blocks, block_size_bytes = 4, mmap.PAGESIZE
with _region(num_blocks=num_blocks, block_size_bytes=block_size_bytes):
total = num_blocks * block_size_bytes
mmap_id = madvise_calls[0][2]
assert madvise_calls == [
(0, mmap.PAGESIZE, mmap_id), # probe
(0, total, mmap_id), # flat pre-fault of the whole region
]
assert fallback_calls == [], (
"native-path constructor must not invoke the fallback helper"
)
def test_madvise_einval_selects_fallback_for_whole_region(monkeypatch):
"""An EINVAL probe must select the fallback for the whole flat region.
This is the regression under test: on Linux < 5.14 the constant is absent
from the `mmap` module, so advice=23 reaches a kernel that rejects it and
every ECSharedRegion construction used to abort with EINVAL.
"""
from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
fallback_calls: list[tuple[int, int]] = []
real_fallback = esr._fallback_populate_write
def _raise_einval(mm, off, ln):
raise OSError(errno.EINVAL, "simulated unsupported kernel")
def _spy_fallback(mm, off, ln):
fallback_calls.append((off, ln))
return real_fallback(mm, off, ln)
monkeypatch.setattr(esr, "_madvise_populate_write", _raise_einval)
monkeypatch.setattr(esr, "_fallback_populate_write", _spy_fallback)
num_blocks, block_size_bytes = 4, mmap.PAGESIZE
with _region(num_blocks=num_blocks, block_size_bytes=block_size_bytes) as r:
assert fallback_calls == [(0, num_blocks * block_size_bytes)]
# Region is still fully usable after taking the fallback path.
r.blocks[0, :4] = torch.tensor([1, 2, 3, 4], dtype=torch.int8)
assert r.blocks[0, :4].tolist() == [1, 2, 3, 4]
def test_non_creator_does_not_probe_or_populate(monkeypatch):
"""Only the creator pre-faults. A second opener must not probe the advice
nor re-touch pages the creator already populated."""
from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
instance_id = str(uuid.uuid4())
r1 = ECSharedRegion(
engine_id=instance_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE
)
try:
calls: list[str] = []
monkeypatch.setattr(
esr,
"_madvise_populate_write",
lambda mm, off, ln: calls.append("madvise"),
)
monkeypatch.setattr(
esr,
"_fallback_populate_write",
lambda mm, off, ln: calls.append("fallback"),
)
r2 = ECSharedRegion(
engine_id=instance_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE
)
try:
assert not r2._is_creator
assert calls == []
finally:
r2.cleanup()
finally:
r1.cleanup()
def test_fallback_populate_write_preserves_bytes_and_faults_pages():
"""Fallback preserves existing bytes and touches every target page.
`|= 0` rather than `= 0` matters here: a peer worker may already have
written EC data into the shared mmap before we pre-fault it.
"""
from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
size = 3 * mmap.PAGESIZE
if not hasattr(mmap, "MADV_DONTNEED"):
pytest.skip("MADV_DONTNEED is unavailable")
mmap_obj = mmap.mmap(
-1,
size,
flags=mmap.MAP_SHARED,
prot=mmap.PROT_READ | mmap.PROT_WRITE,
)
try:
mmap_obj.madvise(mmap.MADV_DONTNEED, 0, size)
if any(_page_residency(mmap_obj, size)):
pytest.skip("kernel did not discard anonymous mmap pages")
mmap_obj[0] = 0xAB
assert _page_residency(mmap_obj, size) == [True, False, False]
esr._fallback_populate_write(mmap_obj, 0, size)
assert _page_residency(mmap_obj, size) == [True, True, True]
assert mmap_obj[0] == 0xAB
finally:
mmap_obj.close()
def test_madvise_unexpected_oserror_propagates(monkeypatch):
"""Only EINVAL triggers the fallback. Other OSErrors (e.g. EIO) must
propagate out of __init__, not be silently masked by the fallback branch.
"""
from vllm.distributed.ec_transfer.ec_connector.cpu import ec_shared_region as esr
monkeypatch.setattr(
esr,
"_madvise_populate_write",
lambda mm, off, ln: (_ for _ in ()).throw(
OSError(errno.EIO, "simulated I/O failure")
),
)
engine_id = str(uuid.uuid4())
try:
with pytest.raises(OSError) as exc_info:
ECSharedRegion(
engine_id=engine_id, num_blocks=4, block_size_bytes=mmap.PAGESIZE
)
assert exc_info.value.errno == errno.EIO
finally:
# __init__ raised past the unlink-owning cleanup(), so drop the
# backing file here rather than leaking it into /dev/shm.
with contextlib.suppress(FileNotFoundError):
os.unlink(f"/dev/shm/vllm_ec_{engine_id}.mmap")
# ── pin_memory ───────────────────────────────────────────────────────────────
def test_pin_memory_success_sets_flag(region):
"""When cudaHostRegister returns 0, _is_pinned flips to True
and cleanup will correspondingly call cudaHostUnregister."""
fake_cudart = MagicMock()
success = MagicMock()
success.value = 0
fake_cudart.cudaHostRegister.return_value = success
fake_cudart.cudaHostUnregister.return_value = success
with (
patch("torch.cuda.is_available", return_value=True),
patch("torch.cuda.cudart", return_value=fake_cudart),
):
region.pin_memory()
assert region._is_pinned is True
# cleanup must pair with cudaHostUnregister exactly once.
region.cleanup()
fake_cudart.cudaHostUnregister.assert_called_once()
def test_pin_memory_failure_leaves_flag_false():
"""If cudaHostRegister fails (non-zero), don't pretend it succeeded —
cleanup must NOT call cudaHostUnregister on memory we never registered."""
r = _make_region()
try:
fake_cudart = MagicMock()
fail = MagicMock()
fail.value = 1 # non-zero == error
fake_cudart.cudaHostRegister.return_value = fail
with (
patch("torch.cuda.is_available", return_value=True),
patch("torch.cuda.cudart", return_value=fake_cudart),
):
r.pin_memory()
assert r._is_pinned is False
# Now run cleanup and verify cudaHostUnregister was NOT called.
r.cleanup()
fake_cudart.cudaHostUnregister.assert_not_called()
finally:
r.cleanup()
def test_pin_memory_noop_without_cuda(region):
"""pin_memory is a no-op when CUDA is not available."""
with patch("torch.cuda.is_available", return_value=False):
region.pin_memory()
assert region._is_pinned is False
# ── cleanup idempotency ───────────────────────────────────────────────────────
def test_cleanup_is_idempotent(region):
region.cleanup()
region.cleanup() # fixture calls a third time — must not raise
tests/v1/ec_connector/unit/test_control.py
-285
View File
@@ -1,285 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for control-layer ABCs, ZmqClientConnection, and ECSession."""
from unittest.mock import MagicMock, patch
import pytest
import zmq
from vllm.distributed.ec_transfer.ec_connector.cpu.control.base import (
ControlConnection,
)
class _ConcreteConn(ControlConnection):
"""Minimal concrete subclass for ABC-instantiation tests."""
def __init__(self):
self._alive = True
self._sent = []
self._inbox = []
@property
def alive(self) -> bool:
return self._alive
def send(self, msg: bytes) -> None:
self._sent.append(msg)
def recv(self) -> list[bytes]:
msgs, self._inbox = self._inbox, []
return msgs
def mark_dead(self) -> None:
self._alive = False
def close(self) -> None:
self._alive = False
def test_control_connection_abc_cannot_be_instantiated_directly():
with pytest.raises(TypeError):
ControlConnection() # type: ignore[abstract]
def test_concrete_connection_send_recv():
conn = _ConcreteConn()
assert conn.alive
conn.send(b"hello")
conn._inbox.append(b"world")
assert conn.recv() == [b"world"]
assert conn._sent == [b"hello"]
def test_mark_dead_sets_alive_false():
conn = _ConcreteConn()
conn.mark_dead()
assert not conn.alive
def test_close_sets_alive_false():
conn = _ConcreteConn()
conn.close()
assert not conn.alive
# ── ZmqClientConnection ───────────────────────────────────────────────────────
from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import ( # noqa: E402
ZmqClientConnection,
ZmqClientTransport,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr # noqa: E402
def _conn(messages: list[list[bytes]], monitor=None) -> ZmqClientConnection:
"""Build a ZmqClientConnection whose dealer yields the given frame-lists."""
dealer = MagicMock()
dealer.recv_multipart.side_effect = [*messages, zmq.Again()]
return ZmqClientConnection(dealer=dealer, monitor=monitor)
def test_zmq_client_connection_recv_extracts_payloads_from_envelopes():
conn = _conn([[b"", b"msg1"], [b"", b"msg2"]])
assert conn.recv() == [b"msg1", b"msg2"]
def test_zmq_client_connection_recv_skips_frames_without_empty_delimiter():
# A frame that does not start with b"" is malformed and must be dropped.
conn = _conn([[b"", b"ok"], [b"no-delimiter", b"bad"]])
assert conn.recv() == [b"ok"]
def test_zmq_client_connection_recv_breaks_on_non_again_exception():
# A socket error should stop the drain and return whatever was collected.
dealer = MagicMock()
dealer.recv_multipart.side_effect = [[b"", b"first"], OSError("broken")]
conn = ZmqClientConnection(dealer=dealer, monitor=None)
assert conn.recv() == [b"first"]
def test_zmq_client_connection_send_prefixes_empty_delimiter():
dealer = MagicMock()
conn = ZmqClientConnection(dealer=dealer, monitor=None)
conn.send(b"payload")
dealer.send_multipart.assert_called_once_with([b"", b"payload"])
def test_zmq_client_connection_mark_dead_does_not_close_socket():
# mark_dead signals liveness only; resource teardown is close()'s job.
dealer = MagicMock()
conn = ZmqClientConnection(dealer=dealer, monitor=None)
conn.mark_dead()
assert not conn.alive
dealer.close.assert_not_called()
def test_zmq_client_connection_close_without_monitor():
dealer = MagicMock()
conn = ZmqClientConnection(dealer=dealer, monitor=None)
conn.close()
dealer.close.assert_called_once_with(linger=0)
assert not conn.alive
def test_zmq_client_connection_close_with_monitor_disables_then_closes_both():
# The disable_monitor → monitor.close → dealer.close sequence must be
# respected: closing the dealer before disabling the monitor can cause
# a zmq assertion in some versions.
dealer = MagicMock()
monitor = MagicMock()
conn = ZmqClientConnection(dealer=dealer, monitor=monitor)
conn.close()
dealer.disable_monitor.assert_called_once()
monitor.close.assert_called_once_with(linger=0)
dealer.close.assert_called_once_with(linger=0)
assert not conn.alive
def test_zmq_client_connection_close_is_idempotent():
dealer = MagicMock()
conn = ZmqClientConnection(dealer=dealer, monitor=None)
conn.close()
conn.close() # must not raise
assert not conn.alive
# ── ZmqClientTransport ────────────────────────────────────────────────────────
def _make_transport() -> tuple[ZmqClientTransport, MagicMock]:
"""Return a ZmqClientTransport with a mocked ZMQ context."""
ctx = MagicMock()
ctx.socket.return_value = MagicMock()
with (
patch(
"vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context",
return_value=ctx,
),
patch(
"vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket",
return_value=MagicMock(),
),
patch(
"vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path",
return_value="tcp://host:0",
),
):
transport = ZmqClientTransport()
transport._ctx = ctx
return transport, ctx
def _inject_conn(transport: ZmqClientTransport, addr: PeerAddr) -> ZmqClientConnection:
"""Inject a ZmqClientConnection directly into the transport pool."""
conn = ZmqClientConnection(dealer=MagicMock(), monitor=MagicMock())
transport._connections[addr] = conn
return conn
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
def test_zmq_client_transport_connect_caches_connection(
mock_ctx_cls, mock_path, mock_sock
):
mock_ctx_cls.return_value = MagicMock()
mock_ctx_cls.return_value.socket.return_value = MagicMock()
mock_sock.return_value = MagicMock()
transport = ZmqClientTransport()
addr: PeerAddr = ("host", 1234)
assert transport.connect(addr) is transport.connect(addr)
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
def test_zmq_client_transport_connect_creates_independent_conns_per_addr(
mock_ctx_cls, mock_path, mock_sock
):
mock_ctx_cls.return_value = MagicMock()
mock_ctx_cls.return_value.socket.return_value = MagicMock()
mock_sock.return_value = MagicMock()
transport = ZmqClientTransport()
conn_a = transport.connect(("host-a", 1111))
conn_b = transport.connect(("host-b", 2222))
assert conn_a is not conn_b
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
def test_zmq_client_transport_poll_dead_detects_disconnected_peer(mock_recv_mon):
transport, _ = _make_transport()
addr: PeerAddr = ("host", 1234)
conn = _inject_conn(transport, addr)
mock_recv_mon.side_effect = [{"event": zmq.EVENT_DISCONNECTED}, zmq.Again()]
dead = transport.poll_dead()
assert addr in dead
assert not conn.alive
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
def test_zmq_client_transport_poll_dead_skips_conn_without_monitor(mock_recv_mon):
transport, _ = _make_transport()
addr: PeerAddr = ("host", 9999)
conn = ZmqClientConnection(dealer=MagicMock(), monitor=None)
transport._connections[addr] = conn
dead = transport.poll_dead()
assert dead == []
mock_recv_mon.assert_not_called()
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.recv_monitor_message")
def test_zmq_client_transport_poll_dead_ignores_non_disconnect_events(mock_recv_mon):
transport, _ = _make_transport()
addr: PeerAddr = ("host", 1234)
conn = _inject_conn(transport, addr)
# A CONNECTED event should not trigger a dead notification.
mock_recv_mon.side_effect = [{"event": zmq.EVENT_CONNECTED}, zmq.Again()]
dead = transport.poll_dead()
assert dead == []
assert conn.alive
def test_zmq_client_transport_remove_closes_and_evicts_conn():
transport, _ = _make_transport()
addr: PeerAddr = ("host", 5678)
conn = _inject_conn(transport, addr)
removed = transport.remove(addr)
assert removed is conn
assert addr not in transport._connections
conn.dealer.close.assert_called()
def test_zmq_client_transport_remove_returns_none_for_unknown_addr():
transport, _ = _make_transport()
assert transport.remove(("ghost", 0)) is None
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_socket")
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.make_zmq_path")
@patch("vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq.zmq.Context")
def test_zmq_client_transport_close_closes_all_connections(
mock_ctx_cls, mock_path, mock_sock
):
mock_ctx = MagicMock()
mock_ctx_cls.return_value = mock_ctx
mock_ctx.socket.return_value = MagicMock()
dealer_a, dealer_b = MagicMock(), MagicMock()
mock_sock.side_effect = [dealer_a, dealer_b]
transport = ZmqClientTransport()
transport.connect(("a", 1))
transport.connect(("b", 2))
transport.close()
dealer_a.close.assert_called()
dealer_b.close.assert_called()
assert transport._connections == {}
tests/v1/ec_connector/unit/test_data.py
-184
View File
@@ -1,184 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for DataTransport ABC and NixlDataTransport."""
from unittest.mock import MagicMock, patch
import pytest
from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import (
DataTransport,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.utils import (
serialize_mem_descriptor,
)
_NIXL_PATH = "vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl"
# ── helpers ───────────────────────────────────────────────────────────────────
def _make_nixl_mock(prep_side_effect=None) -> MagicMock:
nixl = MagicMock()
nixl.get_agent_metadata.return_value = b"agent-meta"
nixl.get_reg_descs.return_value = MagicMock()
nixl.get_xfer_descs.return_value = MagicMock()
nixl.prep_xfer_dlist.side_effect = prep_side_effect or [77, 88]
nixl.add_remote_agent.return_value = "remote-agent-1"
nixl.get_new_notifs.return_value = {}
nixl.check_xfer_state.return_value = "PROC"
return nixl
def _make_transport(nixl_mock=None):
from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
NixlDataTransport,
)
mock = nixl_mock or _make_nixl_mock()
with (
patch(f"{_NIXL_PATH}.NixlWrapper", return_value=mock),
patch(f"{_NIXL_PATH}.nixl_agent_config", return_value=MagicMock()),
):
t = NixlDataTransport("agent", 0, 4, 64, 256)
t._nixl = mock
return t
# ── ABC ───────────────────────────────────────────────────────────────────────
def test_data_transport_abc_cannot_be_instantiated():
with pytest.raises(TypeError):
DataTransport() # type: ignore[abstract]
# ── NixlDataTransport construction ───────────────────────────────────────────
def test_nixl_data_transport_registers_memory_on_init():
nixl = _make_nixl_mock()
_make_transport(nixl)
nixl.register_memory.assert_called_once()
def test_nixl_data_transport_unwinds_registration_if_init_fails():
"""A failure after register_memory must not leave the region pinned.
__init__ propagates, so the caller never gets a transport to deregister.
"""
nixl = _make_nixl_mock(prep_side_effect=RuntimeError("no dlist"))
from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
NixlDataTransport,
)
with (
patch(f"{_NIXL_PATH}.NixlWrapper", return_value=nixl),
patch(f"{_NIXL_PATH}.nixl_agent_config", return_value=MagicMock()),
pytest.raises(RuntimeError, match="no dlist"),
):
NixlDataTransport("agent", 0, 4, 64, 256)
nixl.deregister_memory.assert_called_once_with(nixl.get_reg_descs.return_value)
def test_nixl_data_transport_raises_if_nixl_unavailable():
from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
NixlDataTransport,
)
with (
patch(f"{_NIXL_PATH}.NixlWrapper", None),
patch(f"{_NIXL_PATH}.nixl_agent_config", None),
pytest.raises(RuntimeError, match="requires NIXL"),
):
NixlDataTransport("a", 0, 1, 64, 64)
# ── public accessors ──────────────────────────────────────────────────────────
def test_get_agent_metadata_returns_nixl_blob():
assert _make_transport().get_agent_metadata() == b"agent-meta"
def test_get_mem_descriptor_roundtrips_block_descs():
t = _make_transport()
raw = t.get_mem_descriptor()
assert isinstance(raw, bytes) and len(raw) > 0
# ── add_remote_peer ───────────────────────────────────────────────────────────
def test_add_remote_peer_returns_agent_name_and_stores_handle_internally():
nixl = _make_nixl_mock(prep_side_effect=[77, 88])
t = _make_transport(nixl)
mem_desc = serialize_mem_descriptor([(0, 64, 0)])
agent_name = t.add_remote_peer(b"meta", mem_desc)
assert agent_name == "remote-agent-1"
# Dlist handle stored internally; accessible via _peer_handles.
assert t._peer_handles["remote-agent-1"] == 88
def test_add_remote_peer_registers_remote_agent():
nixl = _make_nixl_mock(prep_side_effect=[77, 88])
t = _make_transport(nixl)
mem_desc = serialize_mem_descriptor([(0, 64, 0)])
t.add_remote_peer(b"fresh-meta", mem_desc)
nixl.add_remote_agent.assert_called_once_with(b"fresh-meta")
# ── remove_remote_peer ────────────────────────────────────────────────────────
def test_remove_remote_peer_delegates_to_nixl():
t = _make_transport()
t.remove_remote_peer("agent-x")
t._nixl.remove_remote_agent.assert_called_once_with("agent-x")
def test_remove_remote_peer_swallows_exceptions():
t = _make_transport()
t._nixl.remove_remote_agent.side_effect = RuntimeError("gone")
t.remove_remote_peer("agent-x") # must not raise
# ── post_read ─────────────────────────────────────────────────────────────────
def test_post_read_raises_on_index_count_mismatch():
nixl = _make_nixl_mock(prep_side_effect=[77, 88])
t = _make_transport(nixl)
mem_desc = serialize_mem_descriptor([(0, 64, 0)])
agent = t.add_remote_peer(b"meta", mem_desc)
with pytest.raises(ValueError, match="mismatch"):
t.post_read([0, 1], agent, [5], notif_msg=b"sid")
def test_post_read_calls_make_prepped_xfer_and_transfer():
nixl = _make_nixl_mock(prep_side_effect=[77, 88])
t = _make_transport(nixl)
mem_desc = serialize_mem_descriptor([(0, 64, 0)])
agent = t.add_remote_peer(b"meta", mem_desc)
nixl.make_prepped_xfer.return_value = 999
handle = t.post_read([0, 1], agent, [5, 6], notif_msg=b"sid")
nixl.make_prepped_xfer.assert_called_once()
nixl.transfer.assert_called_once_with(999)
assert handle == 999
# ── release_xfer_handle / deregister ─────────────────────────────────────────
def test_release_xfer_handle_swallows_exceptions():
t = _make_transport()
t._nixl.release_xfer_handle.side_effect = RuntimeError("stale")
t.release_xfer_handle(42) # must not raise
def test_deregister_swallows_exceptions():
t = _make_transport()
t._nixl.deregister_memory.side_effect = RuntimeError("already gone")
t.deregister() # must not raise
tests/v1/ec_connector/unit/test_ec_connector_metrics.py
-85
View File
@@ -1,85 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the connector-agnostic EC connector metrics plumbing.
Covers only the base infra: no concrete connector type implements these
hooks yet, so every assertion here is about the documented no-op defaults
and the container contracts, mirroring the equivalent KV connector tests.
"""
from unittest.mock import Mock
import pytest
from vllm.config import VllmConfig
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
from vllm.distributed.ec_transfer.ec_connector.example_connector import (
ECExampleConnector,
)
from vllm.distributed.ec_transfer.ec_connector.metrics import (
ECConnectorLogging,
ECConnectorProm,
ECConnectorStats,
)
from vllm.v1.outputs import ECConnectorOutput
pytestmark = pytest.mark.cpu_test
@pytest.fixture
def example_connector(tmp_path):
config = Mock(spec=VllmConfig)
config.ec_transfer_config = Mock()
config.ec_transfer_config.get_from_extra_config = Mock(return_value=str(tmp_path))
config.ec_transfer_config.is_ec_producer = True
config.ec_transfer_config.is_ec_consumer = False
return ECExampleConnector(vllm_config=config, role=ECConnectorRole.SCHEDULER)
def test_base_stats_methods_raise_not_implemented():
"""The base container has no aggregation logic of its own."""
stats = ECConnectorStats()
with pytest.raises(NotImplementedError):
stats.reset()
with pytest.raises(NotImplementedError):
stats.aggregate(stats)
with pytest.raises(NotImplementedError):
stats.reduce()
with pytest.raises(NotImplementedError):
stats.is_empty()
def test_connector_hooks_default_to_none(example_connector):
"""No connector type overrides these yet, so every hook is a no-op."""
assert example_connector.get_ec_connector_stats() is None
assert ECExampleConnector.build_ec_connector_stats() is None
assert (
ECExampleConnector.build_prom_metrics(
vllm_config=Mock(spec=VllmConfig),
metric_types={},
labelnames=[],
per_engine_labelvalues={},
)
is None
)
def test_ec_connector_output_is_empty_reflects_stats_field():
output = ECConnectorOutput()
assert output.is_empty()
output.ec_connector_stats = ECConnectorStats(data={"x": 1})
assert not output.is_empty()
def test_logging_and_prom_are_no_ops_without_ec_transfer_config():
"""A normal run with no EC connector configured must stay silent."""
logging = ECConnectorLogging(ec_transfer_config=None)
logging.log() # Should not raise, nothing to log.
vllm_config = Mock(spec=VllmConfig)
vllm_config.ec_transfer_config = None
prom = ECConnectorProm(vllm_config, labelnames=[], per_engine_labelvalues={})
assert prom.prom_metrics is None
prom.observe({}) # Should not raise.
tests/v1/ec_connector/unit/test_ec_example_connector.py
-747
View File
@@ -1,747 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for ECExampleConnector."""
import os
from unittest.mock import Mock, patch
import pytest
import safetensors
import torch
from vllm.config import VllmConfig
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorRole
from vllm.distributed.ec_transfer.ec_connector.example_connector import (
ECExampleConnector,
ECExampleConnectorMetadata,
MMMeta,
)
from vllm.distributed.ec_transfer.ec_connector.utils import collect_ec_item_metadata
from vllm.multimodal.inputs import MultiModalFeatureSpec, PlaceholderRange
from vllm.v1.core.sched.output import SchedulerOutput
# ------------------ Mock Classes ------------------ #
class MockRequest:
def __init__(self, request_id, mm_hashes: list[str], token_counts: list[int]):
assert len(mm_hashes) == len(token_counts)
self.request_id = request_id
self._token_counts = token_counts
self.mm_features = []
for i, mm_hash in enumerate(mm_hashes):
feature = MultiModalFeatureSpec(
data=None,
modality="image",
identifier=mm_hash,
mm_position=PlaceholderRange(offset=0, length=self._token_counts[i]),
)
self.mm_features.append(feature)
def get_num_encoder_embeds(self, input_id: int) -> int:
assert input_id < len(self._token_counts)
return self._token_counts[input_id]
def test_metadata_preserves_repeated_item_positions():
request = MockRequest("request", ["A", "B", "A"], [4, 4, 4])
resolver = Mock()
resolver.fields_for.return_value = {"image_grid_thw"}
for feature in request.mm_features:
feature.data = Mock()
feature.data.get_data.return_value = {"image_grid_thw": torch.tensor([1, 2, 2])}
items = collect_ec_item_metadata(request.mm_features, resolver)
assert items == {
"A": {"metadata": {"image_grid_thw": [1, 2, 2]}, "item_indices": [0, 2]},
"B": {"metadata": {"image_grid_thw": [1, 2, 2]}, "item_indices": [1]},
}
def test_audio_metadata_survives_processor_cache_hit():
resolver = Mock()
resolver.fields_for.return_value = {"audio_num_tokens"}
feature = MultiModalFeatureSpec(
data=None,
modality="audio",
identifier="audio",
mm_position=PlaceholderRange(
offset=0, length=4, is_embed=torch.tensor([False, True, True, False])
),
)
assert collect_ec_item_metadata([feature], resolver) == {
"audio": {"metadata": {"audio_num_tokens": [2]}, "item_indices": [0]}
}
@pytest.fixture
def temp_storage(tmp_path):
"""Fixture providing temporary storage path."""
return str(tmp_path)
@pytest.fixture
def mock_vllm_config_producer(temp_storage):
"""Fixture providing mock VllmConfig for producer role."""
config = Mock(spec=VllmConfig)
config.ec_transfer_config = Mock()
config.ec_transfer_config.get_from_extra_config = Mock(return_value=temp_storage)
config.ec_transfer_config.is_ec_producer = True
return config
@pytest.fixture
def mock_vllm_config_consumer(temp_storage):
"""Fixture providing mock VllmConfig for consumer role."""
config = Mock(spec=VllmConfig)
config.ec_transfer_config = Mock()
config.ec_transfer_config.get_from_extra_config = Mock(return_value=temp_storage)
config.ec_transfer_config.is_ec_producer = False
return config
@pytest.fixture
def mock_request_with_3_mm():
"""Fixture providing mock Request with 3 multimodal items."""
request_id = "test_req_123"
mm_hashes = ["img_hash_1", "img_hash_2", "img_hash_3"]
token_counts = [100, 150, 200]
request = MockRequest(request_id, mm_hashes, token_counts)
return request
# ------------------ Unit Tests ------------------ #
class TestECExampleConnectorBasics:
"""Test basic EC connector functionality."""
def test_initialization_producer(self, mock_vllm_config_producer, temp_storage):
"""Test connector initializes correctly as producer."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
assert connector.role == ECConnectorRole.SCHEDULER
assert connector.is_producer
assert connector._storage_path == temp_storage
assert connector._mm_datas_need_loads == {}
def test_initialization_consumer(self, mock_vllm_config_consumer, temp_storage):
"""Test connector initializes correctly as consumer."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
assert connector.role == ECConnectorRole.WORKER
assert not connector.is_producer
assert connector._storage_path == temp_storage
def test_role_assignment(self, mock_vllm_config_producer):
"""Test role is correctly assigned."""
scheduler_connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
worker_connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
assert scheduler_connector.role == ECConnectorRole.SCHEDULER
assert worker_connector.role == ECConnectorRole.WORKER
class TestCacheExistence:
"""Test cache existence checking using has_cache_item() API."""
def test_has_cache_item_all_exist_3_items(
self,
mock_vllm_config_producer,
mock_vllm_config_consumer,
mock_request_with_3_mm,
):
"""Test has_cache_item returns True when all 3 caches exist."""
# Test for producer first
producer = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
# Create cache files using save_caches (proper way)
encoder_cache: dict[str, torch.Tensor] = {}
for mm_feature in mock_request_with_3_mm.mm_features:
mm_hash = mm_feature.identifier
encoder_cache[mm_hash] = torch.randn(10, 768)
producer.save_caches(encoder_cache, mm_hash)
# Test using has_cache_item API
producer_result = [
producer.has_cache_item(mm_feature.identifier)
for mm_feature in mock_request_with_3_mm.mm_features
]
# Assert
assert len(producer_result) == 3
assert all(producer_result), f"Expected all True, got {producer_result}"
# Also test consumer can check if cache exists
consumer = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.SCHEDULER,
)
# Test using has_cache_item API
consumer_result = [
consumer.has_cache_item(mm_feature.identifier)
for mm_feature in mock_request_with_3_mm.mm_features
]
# Assert
assert len(consumer_result) == 3
assert all(consumer_result), f"Expected all True, got {consumer_result}"
def test_has_cache_item_none_exist(
self, mock_vllm_config_producer, mock_request_with_3_mm
):
"""Test has_caches returns False when no caches exist."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
# Test without creating any files
result = [
connector.has_cache_item(mm_feature.identifier)
for mm_feature in mock_request_with_3_mm.mm_features
]
# Assert
assert len(result) == 3
assert not any(result), f"Expected all False, got {result}"
def test_has_cache_item_partial_exist(
self, mock_vllm_config_producer, mock_request_with_3_mm
):
"""Test has_caches with some caches existing (1 of 3)."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
# Create only the second cache file
mm_hash_second = mock_request_with_3_mm.mm_features[1].identifier
encoder_cache = {mm_hash_second: torch.randn(10, 768)}
connector.save_caches(encoder_cache, mm_hash_second)
# Test
result = [
connector.has_cache_item(mm_feature.identifier)
for mm_feature in mock_request_with_3_mm.mm_features
]
# Assert
assert len(result) == 3
assert not result[0] # First doesn't exist
assert result[1] # Second exists
assert not result[2] # Third doesn't exist
def test_has_cache_item_miss_creates_no_folder(
self, mock_vllm_config_consumer, temp_storage
):
"""Test a cache miss leaves no empty folder in the shared storage."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.SCHEDULER,
)
assert not connector.has_cache_item("missing_hash")
assert os.listdir(temp_storage) == []
class TestStateManagement:
"""Test connector state management."""
def test_update_state_after_alloc_3_items(
self, mock_vllm_config_producer, mock_request_with_3_mm
):
"""Test state update after allocation for 3 MM items."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
# Initial state should be empty
assert len(connector._mm_datas_need_loads) == 0
# Update state for all 3 items (mock cache existence)
with patch.object(connector, "has_cache_item", return_value=True):
for i in range(3):
connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
# Check state updated for all 3
assert len(connector._mm_datas_need_loads) == 3
assert "img_hash_1" in connector._mm_datas_need_loads
assert "img_hash_2" in connector._mm_datas_need_loads
assert "img_hash_3" in connector._mm_datas_need_loads
assert connector._mm_datas_need_loads["img_hash_1"] == 100
assert connector._mm_datas_need_loads["img_hash_2"] == 150
assert connector._mm_datas_need_loads["img_hash_3"] == 200
def test_build_connector_meta_3_items(
self, mock_vllm_config_producer, mock_request_with_3_mm
):
"""Test metadata building for 3 MM items."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
# Setup state for all 3 items (mock cache existence)
with patch.object(connector, "has_cache_item", return_value=True):
for i in range(3):
connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
# Build metadata
scheduler_output = Mock(spec=SchedulerOutput)
metadata = connector.build_connector_meta(scheduler_output)
# Assert
assert isinstance(metadata, ECExampleConnectorMetadata)
assert len(metadata.mm_datas) == 3
assert metadata.mm_datas[0].mm_hash == "img_hash_1"
assert metadata.mm_datas[0].num_token == 100
assert metadata.mm_datas[1].mm_hash == "img_hash_2"
assert metadata.mm_datas[1].num_token == 150
assert metadata.mm_datas[2].mm_hash == "img_hash_3"
assert metadata.mm_datas[2].num_token == 200
# State should be cleared after building
assert len(connector._mm_datas_need_loads) == 0
def test_build_connector_meta_empty(self, mock_vllm_config_producer):
"""Test metadata building with empty state."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
scheduler_output = Mock(spec=SchedulerOutput)
metadata = connector.build_connector_meta(scheduler_output)
assert isinstance(metadata, ECExampleConnectorMetadata)
assert len(metadata.mm_datas) == 0
def test_state_cleared_after_metadata_build(
self, mock_vllm_config_producer, mock_request_with_3_mm
):
"""Test that state is properly cleared after building metadata."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
# Add state (mock cache existence)
with patch.object(connector, "has_cache_item", return_value=True):
for i in range(3):
connector.update_state_after_alloc(mock_request_with_3_mm, index=i)
assert len(connector._mm_datas_need_loads) == 3
# Build metadata (should clear state)
scheduler_output = Mock(spec=SchedulerOutput)
connector.build_connector_meta(scheduler_output)
# State should be empty
assert len(connector._mm_datas_need_loads) == 0
# Build again should return empty metadata
metadata2 = connector.build_connector_meta(scheduler_output)
assert len(metadata2.mm_datas) == 0
class TestCacheSaving:
"""Test encoder cache saving (producer only)."""
def test_save_caches_producer_3_items(
self, mock_vllm_config_producer, mock_request_with_3_mm, temp_storage
):
"""Test cache saving as producer for 3 different MM items."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
# Create and save 3 different caches
mm_hashes = [f.identifier for f in mock_request_with_3_mm.mm_features]
encoder_cache: dict[str, torch.Tensor] = {}
for mm_hash in mm_hashes:
encoder_cache[mm_hash] = torch.randn(10, 768)
connector.save_caches(encoder_cache, mm_hash)
# Verify all files exist using has_cache_item
result = [
connector.has_cache_item(mm_feature.identifier)
for mm_feature in mock_request_with_3_mm.mm_features
]
assert all(result), f"Not all caches were saved: {result}"
# Verify each file's content
for mm_hash in mm_hashes:
filename = connector._generate_filename_debug(mm_hash)
loaded = safetensors.torch.load_file(filename)
assert "ec_cache" in loaded
assert torch.allclose(loaded["ec_cache"], encoder_cache[mm_hash].cpu())
def test_save_caches_consumer_skips(self, mock_vllm_config_consumer):
"""Test cache saving is skipped for consumer."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
mm_hash = "test_hash_consumer"
encoder_cache = {mm_hash: torch.randn(10, 768)}
# Save should not raise but also not create file
connector.save_caches(encoder_cache, mm_hash)
# Verify file doesn't exist using has_cache_item
result = connector.has_cache_item(mm_hash)
assert not result, "Consumer should not save caches"
class TestCacheLoading:
"""Test encoder cache loading (consumer)."""
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
def test_start_load_caches_consumer_3_items(
self,
mock_vllm_config_producer,
mock_vllm_config_consumer,
mock_request_with_3_mm,
temp_storage,
):
"""Test consumer loads 3 caches from storage."""
# First, create producer to save caches
producer = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
# Producer saves 3 caches
mm_hashes = [f.identifier for f in mock_request_with_3_mm.mm_features]
saved_caches = {}
for mm_hash in mm_hashes:
saved_caches[mm_hash] = torch.randn(10, 768)
producer.save_caches(saved_caches, mm_hash)
# Now consumer loads
consumer = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
# Setup metadata for all 3
metadata = ECExampleConnectorMetadata()
for mm_hash in mm_hashes:
metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
consumer.bind_connector_metadata(metadata)
# Load
encoder_cache: dict[str, torch.Tensor] = {}
consumer.start_load_caches(encoder_cache=encoder_cache)
# Verify all 3 loaded
assert len(encoder_cache) == 3
for mm_hash in mm_hashes:
assert mm_hash in encoder_cache, f"{mm_hash} missing in encoder_cache"
assert encoder_cache[mm_hash].is_cuda, (
f"{mm_hash} cache is in {encoder_cache[mm_hash].device}"
)
assert torch.allclose(
encoder_cache[mm_hash].cpu(), saved_caches[mm_hash]
), f"{mm_hash} cache saved and loaded tesnor are not the same"
def test_start_load_caches_skip_existing(
self, mock_vllm_config_producer, mock_vllm_config_consumer, temp_storage
):
"""Test cache loading skips already cached items."""
# Setup: producer saves cache
producer = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
mm_hash = "existing_hash"
saved_cache = torch.randn(10, 768)
producer.save_caches({mm_hash: saved_cache}, mm_hash)
# Consumer setup
consumer = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
metadata = ECExampleConnectorMetadata()
metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
consumer.bind_connector_metadata(metadata)
# Pre-populate encoder_cache with different value
existing_cache = torch.randn(5, 512)
encoder_cache = {mm_hash: existing_cache}
# Load (should skip since already exists)
with patch("safetensors.torch.load_file") as mock_load:
consumer.start_load_caches(encoder_cache=encoder_cache)
# Should not call load_file since cache exists
mock_load.assert_not_called()
# Verify original cache unchanged
assert torch.equal(encoder_cache[mm_hash], existing_cache)
def test_start_load_caches_empty_metadata(self, mock_vllm_config_consumer):
"""Test loading with empty metadata does nothing."""
consumer = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
# Setup empty metadata
metadata = ECExampleConnectorMetadata()
consumer.bind_connector_metadata(metadata)
# Load (should not raise)
encoder_cache: dict[str, torch.Tensor] = {}
consumer.start_load_caches(encoder_cache=encoder_cache)
# Cache should remain empty
assert len(encoder_cache) == 0
def test_start_load_caches_pins_device_to_current_rank(
self, mock_vllm_config_producer, mock_vllm_config_consumer
):
"""Load must pin to the selected device under tensor parallelism."""
# Produce a real cache file so the load path has something to open.
producer = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
mm_hash = "rank_pinning_hash"
producer.save_caches({mm_hash: torch.zeros(4, 8)}, mm_hash)
consumer = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
metadata = ECExampleConnectorMetadata()
metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
consumer.bind_connector_metadata(metadata)
mock_platform = Mock()
mock_platform.device_type = "cuda"
mock_platform.is_cuda_alike.return_value = True
mock_platform.current_device.return_value = 2
# Simulate running on TP rank 2 without requiring a multi-GPU host.
captured: dict = {}
def fake_load_file(filename, device):
captured["device"] = device
return {"ec_cache": torch.zeros(4, 8)}
with (
patch("vllm.platforms.current_platform", mock_platform),
patch("safetensors.torch.load_file", side_effect=fake_load_file),
):
encoder_cache: dict[str, torch.Tensor] = {}
consumer.start_load_caches(encoder_cache=encoder_cache)
assert captured["device"] == "cuda:2", (
f"Expected device='cuda:2' (pinned to current rank), "
f"got {captured['device']!r}. A bare 'cuda' string makes "
"safetensors load onto cuda:0, which breaks non-rank-0 TP workers."
)
def test_start_load_caches_non_cuda_platform_uses_device_type(
self, mock_vllm_config_producer, mock_vllm_config_consumer
):
"""Non-CUDA platforms pass `device_type` through unchanged (no rank suffix)."""
producer = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
mm_hash = "cpu_device_hash"
producer.save_caches({mm_hash: torch.zeros(2, 4)}, mm_hash)
consumer = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
metadata = ECExampleConnectorMetadata()
metadata.add_mm_data(MMMeta.make_meta(mm_hash, 100))
consumer.bind_connector_metadata(metadata)
mock_platform = Mock()
mock_platform.device_type = "cpu"
mock_platform.is_cuda_alike.return_value = False
captured: dict = {}
def fake_load_file(filename, device):
captured["device"] = device
return {"ec_cache": torch.zeros(2, 4)}
with (
patch("vllm.platforms.current_platform", mock_platform),
patch("safetensors.torch.load_file", side_effect=fake_load_file),
):
encoder_cache: dict[str, torch.Tensor] = {}
consumer.start_load_caches(encoder_cache=encoder_cache)
assert captured["device"] == "cpu"
class TestFilenameGeneration:
"""Test filename and path generation."""
def test_generate_foldername(self, mock_vllm_config_producer, temp_storage):
"""Test folder name generation."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
mm_hash = "test_folder_hash"
folder = connector._generate_foldername_debug(mm_hash)
assert folder == os.path.join(temp_storage, mm_hash)
assert os.path.isdir(folder) # Should be created
def test_generate_filename(self, mock_vllm_config_producer, temp_storage):
"""Test filename generation."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
mm_hash = "test_file_hash"
filename = connector._generate_filename_debug(mm_hash)
expected = os.path.join(temp_storage, mm_hash, "encoder_cache.safetensors")
assert filename == expected
assert os.path.isdir(os.path.dirname(filename)) # Folder created
def test_generate_filename_consistency(self, mock_vllm_config_producer):
"""Test filename generation is consistent."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
mm_hash = "consistency_hash"
filename1 = connector._generate_filename_debug(mm_hash)
filename2 = connector._generate_filename_debug(mm_hash)
assert filename1 == filename2
class TestMetadataBindingLifecycle:
"""Test metadata binding and clearing lifecycle."""
def test_bind_connector_metadata(self, mock_vllm_config_consumer):
"""Test binding connector metadata."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
metadata = ECExampleConnectorMetadata()
metadata.add_mm_data(MMMeta.make_meta("hash_1", 100))
connector.bind_connector_metadata(metadata)
assert connector._connector_metadata is metadata
def test_clear_connector_metadata(self, mock_vllm_config_consumer):
"""Test clearing connector metadata."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
metadata = ECExampleConnectorMetadata()
connector.bind_connector_metadata(metadata)
connector.clear_connector_metadata()
assert connector._connector_metadata is None
def test_get_connector_metadata(self, mock_vllm_config_consumer):
"""Test getting connector metadata."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
metadata = ECExampleConnectorMetadata()
connector.bind_connector_metadata(metadata)
retrieved = connector._get_connector_metadata()
assert retrieved is metadata
def test_get_connector_metadata_not_set(self, mock_vllm_config_consumer):
"""Test getting metadata when not set raises."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
with pytest.raises(AssertionError):
connector._get_connector_metadata()
class TestEdgeCases:
"""Test edge cases and error handling."""
def test_save_empty_cache(self, mock_vllm_config_producer):
"""Test saving empty tensor."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.WORKER,
)
mm_hash = "empty_hash"
encoder_cache = {mm_hash: torch.empty(0)}
# Should not raise
connector.save_caches(encoder_cache, mm_hash)
def test_load_nonexistent_cache(self, mock_vllm_config_consumer):
"""Test loading cache that doesn't exist raises error."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_consumer,
role=ECConnectorRole.WORKER,
)
metadata = ECExampleConnectorMetadata()
metadata.add_mm_data(MMMeta.make_meta("nonexistent_hash", 100))
connector.bind_connector_metadata(metadata)
encoder_cache: dict[str, torch.Tensor] = {}
# Should raise FileNotFoundError
with pytest.raises(FileNotFoundError):
connector.start_load_caches(encoder_cache=encoder_cache)
def test_has_cache_item_empty_request(self, mock_vllm_config_producer):
"""Test has_cache_item with a nonexistent identifier."""
connector = ECExampleConnector(
vllm_config=mock_vllm_config_producer,
role=ECConnectorRole.SCHEDULER,
)
result = connector.has_cache_item("nonexistent_hash")
assert result is False
tests/v1/ec_connector/unit/test_ec_output_aggregator.py
-131
View File
@@ -1,131 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for ECOutputAggregator."""
import pytest
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorWorkerMetadata
from vllm.distributed.ec_transfer.ec_connector.metrics import ECConnectorStats
from vllm.distributed.ec_transfer.ec_connector.utils import ECOutputAggregator
from vllm.distributed.kv_transfer.kv_connector.utils import KVOutputAggregator
from vllm.v1.outputs import (
EMPTY_MODEL_RUNNER_OUTPUT,
ECConnectorOutput,
KVConnectorOutput,
ModelRunnerOutput,
)
pytestmark = pytest.mark.cpu_test
class FakeWorkerMeta(ECConnectorWorkerMetadata):
"""Records merge order. `aggregate` returns a new object, as the base class
declares: an aggregator discarding the return value would lose the merge.
"""
def __init__(self, saves: list[str]):
self.saves = saves
def aggregate(self, other: "FakeWorkerMeta") -> "FakeWorkerMeta":
return FakeWorkerMeta(self.saves + other.saves)
class FakeConnectorStats(ECConnectorStats):
"""Records merge order, mirroring FakeWorkerMeta above."""
def __init__(self, saves: list[str]):
self.saves = saves
def aggregate(self, other: "FakeConnectorStats") -> "FakeConnectorStats":
return FakeConnectorStats(self.saves + other.saves)
def _worker_output(ec_output: ECConnectorOutput | None) -> ModelRunnerOutput:
return ModelRunnerOutput(
req_ids=[], req_id_to_index={}, ec_connector_output=ec_output
)
def test_aggregate_folds_every_rank_onto_output_rank():
"""EC work done on any rank reaches the scheduler via output_rank's output.
The middle rank reports no worker metadata or stats: it must neither seed
nor clobber either accumulator.
"""
outputs = [
_worker_output(
ECConnectorOutput(
finished_sending={"mm0"},
ec_connector_stats=FakeConnectorStats(["mm0"]),
ec_connector_worker_meta=FakeWorkerMeta(["mm0"]),
)
),
_worker_output(ECConnectorOutput(finished_recving={"mm1"})),
_worker_output(
ECConnectorOutput(
ec_connector_stats=FakeConnectorStats(["mm2"]),
ec_connector_worker_meta=FakeWorkerMeta(["mm2"]),
)
),
]
result = ECOutputAggregator().aggregate(outputs, output_rank=2)
assert result is outputs[2]
assert result.ec_connector_output.finished_sending == {"mm0"}
assert result.ec_connector_output.finished_recving == {"mm1"}
assert result.ec_connector_output.ec_connector_stats.saves == ["mm0", "mm2"]
assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0", "mm2"]
def test_aggregate_leaves_no_ec_output_when_no_worker_reported():
"""Empty per-worker reports must not reach the scheduler as an empty object."""
outputs = [_worker_output(ECConnectorOutput()), _worker_output(ECConnectorOutput())]
result = ECOutputAggregator().aggregate(outputs, output_rank=0)
assert result is outputs[0]
assert result.ec_connector_output is None
assert ECOutputAggregator().aggregate([None], output_rank=0) is None
def test_aggregate_does_not_write_through_the_shared_empty_output():
"""A rank with nothing to report yields the shared empty output singleton.
Folding another rank's metadata onto it must not write through to the
module-level object, which every later step would then carry.
"""
outputs = [
_worker_output(
ECConnectorOutput(ec_connector_worker_meta=FakeWorkerMeta(["mm0"]))
),
EMPTY_MODEL_RUNNER_OUTPUT,
]
result = ECOutputAggregator().aggregate(outputs, output_rank=1)
assert EMPTY_MODEL_RUNNER_OUTPUT.ec_connector_output is None
assert result is not EMPTY_MODEL_RUNNER_OUTPUT
assert result.ec_connector_output.ec_connector_worker_meta.saves == ["mm0"]
def test_chaining_with_kv_aggregator_preserves_both_outputs():
"""MultiprocExecutor chains both aggregators and keeps only the last result,
so each must merge onto the same output_rank output rather than replace it.
"""
outputs = [
_worker_output(ECConnectorOutput(finished_sending={"mm0"})),
_worker_output(None),
]
outputs[1].kv_connector_output = KVConnectorOutput(invalid_block_ids={7})
result = None
for aggregator in (
KVOutputAggregator(expected_finished_count=1),
ECOutputAggregator(),
):
result = aggregator.aggregate(outputs, output_rank=1)
assert result is outputs[1]
assert result.kv_connector_output.invalid_block_ids == {7}
assert result.ec_connector_output.finished_sending == {"mm0"}
tests/v1/ec_connector/unit/test_ec_transfer_params.py
-129
View File
@@ -1,129 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Sanity tests for ec_transfer_params protocol plumbing.
No running engine required.
"""
from unittest.mock import MagicMock
import pytest
from vllm.entrypoints.openai.chat_completion.protocol import (
ChatCompletionRequest,
)
from vllm.outputs import CompletionOutput, RequestOutput
from vllm.sampling_params import SamplingParams
from vllm.v1.request import Request, RequestStatus
# tests.v1.core.utils transitively imports the example KV connector, which
# requires the compiled CUDA flash-attention extensions; skip cleanly rather
# than erroring at collection when those extensions aren't built.
_core_utils = pytest.importorskip(
"tests.v1.core.utils",
reason="requires compiled CUDA flash-attention extensions",
exc_type=ImportError,
)
create_scheduler = _core_utils.create_scheduler
EC_PARAMS: dict = {"mm_hash_abc": {"peer_host": "10.0.0.1", "peer_port": 5501}}
def test_ec_transfer_params_routed_to_sampling_params_extra_args():
"""ec_transfer_params on the request must land in SamplingParams.extra_args."""
req = ChatCompletionRequest(
model="test-model",
messages=[{"role": "user", "content": "hi"}],
max_tokens=5,
ec_transfer_params=EC_PARAMS,
)
sp = req.to_sampling_params(max_tokens=5, default_sampling_params={})
assert sp.extra_args is not None
assert sp.extra_args.get("ec_transfer_params") == EC_PARAMS
def test_request_output_add_propagates_ec_transfer_params():
"""RequestOutput.add() must carry ec_transfer_params forward to the caller."""
def _out(ec_params):
return RequestOutput(
request_id="r1",
prompt="p",
prompt_token_ids=[1],
prompt_logprobs=None,
outputs=[
CompletionOutput(
index=0,
text="",
token_ids=[],
cumulative_logprob=0.0,
logprobs=None,
finish_reason=None,
)
],
finished=False,
ec_transfer_params=ec_params,
)
accumulated = _out(None)
accumulated.add(_out(EC_PARAMS), aggregate=True)
assert accumulated.ec_transfer_params == EC_PARAMS
def test_request_reads_ec_transfer_params_from_extra_args():
"""v1 Request must pull ec_transfer_params out of SamplingParams.extra_args."""
sp = SamplingParams(extra_args={"ec_transfer_params": EC_PARAMS})
req = Request(
request_id="r1",
prompt_token_ids=[1, 2, 3],
sampling_params=sp,
pooling_params=None,
)
assert req.ec_transfer_params == EC_PARAMS
def test_free_request_calls_ec_connector_and_surfaces_params():
"""_free_request must call ec_connector.request_finished() and return its params."""
sp = SamplingParams(max_tokens=1)
sp.update_from_generation_config({}, 50256)
request = Request(
request_id="test-req",
prompt_token_ids=[1, 2, 3],
sampling_params=sp,
pooling_params=None,
client_index=0,
)
scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer")
scheduler.add_request(request)
request.status = RequestStatus.FINISHED_STOPPED
mock_ec = MagicMock()
mock_ec.request_finished.return_value = (False, EC_PARAMS)
scheduler.ec_connector = mock_ec
kv_params, ec_params = scheduler._free_request(request)
mock_ec.request_finished.assert_called_once_with(request)
assert ec_params == EC_PARAMS
assert kv_params is None
def test_free_request_without_ec_connector_returns_none():
"""When no EC connector is configured, ec_transfer_params must be None."""
sp = SamplingParams(max_tokens=1)
sp.update_from_generation_config({}, 50256)
request = Request(
request_id="test-req",
prompt_token_ids=[1, 2, 3],
sampling_params=sp,
pooling_params=None,
client_index=0,
)
scheduler = create_scheduler(use_ec_connector=True, ec_role="ec_producer")
scheduler.add_request(request)
request.status = RequestStatus.FINISHED_STOPPED
kv_params, ec_params = scheduler._free_request(request)
assert ec_params is None
assert kv_params is None
tests/v1/ec_connector/unit/test_epd_proxy_round_robin.py
-133
View File
@@ -1,133 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Encoder fan-out fairness in the disaggregated EPD proxy.
Exercises the REAL ``encoder_rr_assignment`` helper loaded from its
``examples/`` path -- no routing logic is re-implemented here, so a future
change to that function is what these tests exercise.
Regression target: the fan-out cursor used to restart at e_urls[0] on every
incoming request (``e_urls[i % len(e_urls)] for i in range(len(mm_items))``),
so single-item requests -- the common case -- always hit the first encoder
instance and left the rest idle. ``encoder_rr_assignment`` threads a cursor
across calls so consecutive requests keep advancing through every URL.
"""
import importlib.util
import sys
from collections import Counter
from pathlib import Path
import pytest
PROXY_REL = "examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py"
def _load_proxy_module():
path = Path(__file__).parents[4] / PROXY_REL
spec = importlib.util.spec_from_file_location("disagg_epd_proxy_under_test", path)
assert spec is not None and spec.loader is not None
module = importlib.util.module_from_spec(spec)
sys.modules[spec.name] = module
spec.loader.exec_module(module)
return module
@pytest.fixture(scope="module")
def proxy():
return _load_proxy_module()
@pytest.fixture(scope="module")
def assign(proxy):
return proxy.encoder_rr_assignment
def _drive(assign, e_urls, counts):
"""Feed a sequence of request item-counts through the cursor, as
fanout_encoder_primer does one request at a time."""
cursor = 0
all_urls = []
for count in counts:
urls, cursor = assign(e_urls, cursor, count)
assert len(urls) == count
all_urls.append(urls)
return all_urls
@pytest.mark.parametrize("n_urls", [1, 2, 3, 5])
def test_full_url_space_is_covered_uniformly(assign, n_urls):
e_urls = [f"E{i}" for i in range(n_urls)]
# Three full cycles of single-item requests.
urls = _drive(assign, e_urls, counts=[1] * (n_urls * 3))
hits = Counter(u for req in urls for u in req)
assert set(hits) == set(e_urls)
assert max(hits.values()) == min(hits.values())
def test_single_item_requests_rotate_through_all_encoders(assign):
# Exact scenario from the bug report: 3 encoders, one image/request.
e_urls = ["E0", "E1", "E2"]
routed = _drive(assign, e_urls, counts=[1] * 6)
assert [urls[0] for urls in routed] == ["E0", "E1", "E2", "E0", "E1", "E2"]
def test_cursor_stays_contiguous_across_varying_item_counts(assign):
e_urls = ["E0", "E1", "E2"]
routed = _drive(assign, e_urls, counts=[2, 1, 3, 1])
assert routed == [
["E0", "E1"],
["E2"],
["E0", "E1", "E2"],
["E0"],
]
def test_single_encoder_always_resolves_to_it(assign):
urls, next_cursor = assign(["E0"], 0, 4)
assert urls == ["E0"] * 4
assert next_cursor == 0
def test_generic_e_p_d_routing_without_mooncake_is_supported(proxy):
proxy.validate_ec_consumer_routing(["http://prefill"], [])
def test_mooncake_e_pd_routing_is_supported(proxy):
proxy.validate_ec_consumer_routing([], ["tcp://decode:19019"])
def test_mooncake_independent_prefill_routing_fails_fast(proxy):
with pytest.raises(ValueError, match=r"supports E\+PD only"):
proxy.validate_ec_consumer_routing(["http://prefill"], ["tcp://decode:19019"])
def test_decode_rewrite_preserves_engine_reported_ec_hash(proxy):
request = {
"messages": [
{
"role": "user",
"content": [{"type": "image_url", "image_url": {"url": "image"}}],
}
]
}
rewritten = proxy.rewrite_for_decode(
request,
{
0: {
"mm_hash": "proxy-uuid",
"ec_mm_hash": "engine-derived-hash",
"transfer_id": "transfer",
"image_grid_thw": [1, 2, 3],
}
},
)
assert rewritten["messages"][0]["content"][0]["uuid"] == "proxy-uuid"
assert rewritten["messages"][0]["content"][0]["image_embeds"] == {
"image_grid_thw": [1, 2, 3]
}
assert rewritten["ec_transfer_params"]["ec_items"] == [
{"mm_hash": "engine-derived-hash", "transfer_id": "transfer"}
]
tests/v1/ec_connector/unit/test_metadata.py
-96
View File
@@ -1,96 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for ECCPUConnectorMetadata."""
from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUConnectorMetadata,
)
def test_metadata_is_ec_connector_metadata():
"""ECCPUConnectorMetadata must be a subclass of ECConnectorMetadata."""
meta = ECCPUConnectorMetadata()
assert isinstance(meta, ECConnectorMetadata)
def test_metadata_initialization():
"""ECCPUConnectorMetadata initializes with empty saves and loads dicts."""
meta = ECCPUConnectorMetadata()
assert hasattr(meta, "saves")
assert hasattr(meta, "loads")
assert meta.saves == {}
assert meta.loads == {}
def test_metadata_saves_dict_operations():
"""The saves dict supports standard dict operations."""
meta = ECCPUConnectorMetadata()
# Add multiple entries
meta.saves["mm_hash_1"] = [0, 1, 2]
meta.saves["mm_hash_2"] = [10, 11]
assert "mm_hash_1" in meta.saves
assert "mm_hash_2" in meta.saves
assert len(meta.saves) == 2
assert meta.saves["mm_hash_1"] == [0, 1, 2]
assert meta.saves["mm_hash_2"] == [10, 11]
def test_metadata_loads_carry_a_transfer_id_and_blocks():
"""Each load entry is a (transfer_id, block_ids) pair, so the worker can
report the transfer rather than the mm_hash."""
meta = ECCPUConnectorMetadata()
meta.loads["mm_hash_1"] = (0, [5, 6, 7])
meta.loads["mm_hash_2"] = (1, [100])
assert len(meta.loads) == 2
assert meta.loads["mm_hash_1"] == (0, [5, 6, 7])
transfer_id, block_ids = meta.loads["mm_hash_2"]
assert transfer_id == 1
assert block_ids == [100]
def test_metadata_saves_and_loads_are_independent():
"""Saves and loads dicts are independent."""
meta = ECCPUConnectorMetadata()
meta.saves["key"] = [1, 2]
meta.loads["key"] = (0, [3, 4])
# Same key but different dicts should not interfere
assert meta.saves["key"] == [1, 2]
assert meta.loads["key"] == (0, [3, 4])
# Modifying one should not affect the other
meta.saves["other"] = [99]
assert "other" not in meta.loads
def test_metadata_block_indices_are_integers():
"""Block indices in saves/loads lists should be integers."""
meta = ECCPUConnectorMetadata()
meta.saves["hash"] = [0, 100, 999, 65535]
meta.loads["hash"] = (0, [1, 2, 3])
for idx in meta.saves["hash"]:
assert isinstance(idx, int)
for idx in meta.loads["hash"][1]:
assert isinstance(idx, int)
def test_metadata_multiple_instances_are_independent():
"""Different ECCPUConnectorMetadata instances should not share state."""
meta1 = ECCPUConnectorMetadata()
meta2 = ECCPUConnectorMetadata()
meta1.saves["x"] = [1]
meta2.saves["y"] = [2]
assert "x" in meta1.saves
assert "x" not in meta2.saves
assert "y" not in meta1.saves
assert "y" in meta2.saves
tests/v1/ec_connector/unit/test_protocol.py
-41
View File
@@ -1,41 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import msgspec
from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import (
EC_CONNECTOR_VERSION,
XferAck,
XferReq,
XferStatus,
compute_ec_compatibility_hash,
)
def test_xferreq_roundtrip():
req = XferReq(mm_hash="h1", compatibility_hash="c1", session_id="s1")
data = msgspec.msgpack.encode(req)
out = msgspec.msgpack.decode(data, type=XferReq)
assert out.mm_hash == "h1"
assert out.compatibility_hash == "c1"
assert out.connector_version == EC_CONNECTOR_VERSION
def test_xferack_roundtrip_ok():
ack = XferAck(
mm_hash="h1",
status=XferStatus.OK,
src_block_indices=[1, 2],
agent_metadata=b"meta",
mem_descriptor=b"desc",
)
out = msgspec.msgpack.decode(msgspec.msgpack.encode(ack), type=XferAck)
assert out.status == XferStatus.OK
assert out.src_block_indices == [1, 2]
def test_compat_hash_deterministic_and_sensitive():
a = compute_ec_compatibility_hash("0.1", "m", "float16", 64)
b = compute_ec_compatibility_hash("0.1", "m", "float16", 64)
c = compute_ec_compatibility_hash("0.1", "m", "float16", 128)
assert a == b
assert a != c
tests/v1/ec_connector/unit/test_scheduler_nixl_consumer.py
-461
View File
@@ -1,461 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import time
import uuid
import pytest
import torch
import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
ECSharedRegion,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
from vllm.multimodal.inputs import (
MultiModalFieldElem,
MultiModalKwargsItem,
MultiModalSharedField,
)
_N, _BS, _HID, _ES = 16, 64, 32, 2
class _Pos:
def __init__(self, offset, length):
self.offset, self.length = offset, length
def get_num_embeds(self):
return self.length
class _Feature:
def __init__(self, mm_hash, length=1, data=None):
self.mm_hash = mm_hash
self.identifier = mm_hash
self.mm_position = _Pos(0, length)
self.data = data
self.modality = "image"
class _Request:
def __init__(self, features, params=None, req_id="r1"):
self.mm_features = features
self.request_id = req_id
self.ec_transfer_params = params
class _FakeResults:
def __init__(self):
self.completed = set()
self.tombstoned = set()
self.quarantined = set()
self.retryable = set()
self.settled = []
class _FakeSession:
"""Records started transfers; lets the test flip them to completed."""
def __init__(self):
self.started = []
self._results = _FakeResults()
def start_xfer(self, mm_hash, block_indices, deadline):
self.started.append(mm_hash)
def poll(self, messages, now):
pass
def take_results(self):
r, self._results = self._results, _FakeResults()
return r
def close(self):
pass
class _FakeTransport:
"""Stand-in for ZmqClientTransport so _poll_step runs without ZMQ."""
def poll(self):
return {}
def poll_dead(self):
return []
def close(self):
pass
_MISSING = object()
def _consumer_sched(monkeypatch):
def _region(cfg):
return ECSharedRegion(
engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS
)
monkeypatch.setattr(sched_mod, "create_ec_shared_region", _region)
# ec_enable_nixl defaults to False, so this builds gate-off; the fields
# below flip on the NIXL consumer state directly.
# No measured widths: a consumer takes each shape from the announcement.
s = ECCPUScheduler(
create_ec_vllm_config(ec_role="ec_consumer", encoder_output_widths={})
)
# Turn on NIXL consumer state without constructing real transports.
s._nixl_enabled = True
s._transport = _FakeTransport()
s._data = None
s._compat_hash = "c"
s._metadata_resolver._cache["image"] = {"image_grid_thw"}
return s
def _params(mm_hash, length):
return {
mm_hash: {
"peer_host": "h",
"peer_port": 1,
"size_bytes": length * _HID * _ES,
"shape": [length, _HID],
}
}
def test_new_remote_read_defers_then_completes(monkeypatch):
"""The read is sized from the announced shape, not a local width."""
width = 4 * _HID
s = _consumer_sched(monkeypatch)
fake = _FakeSession()
# Route _start_xfer to our fake session instead of real ZMQ/NIXL, but
# still reserve a real not-ready cache entry so mark_ready works.
def _fake_start(mm_hash, info, size, shape):
assert size == width * _ES
assert shape == (1, width)
entry = s._cache.alloc(mm_hash, size // _BS, shape)
assert entry is not None
fake.started.append(mm_hash)
return True
monkeypatch.setattr(s, "_start_xfer", _fake_start)
params = _params("h1", 1)
params["h1"]["size_bytes"] = width * _ES
params["h1"]["shape"] = [1, width]
req = _Request([_Feature("h1", 1)], params=params)
# Step 1: unseen remote item -> read started, request deferred.
assert s.ensure_cache_available(req, 0) is False
assert "h1" in s._in_flight
assert "h1" in fake.started
s.build_connector_meta(scheduler_output=None) # re-arms _first_in_batch
# Session reports completion this step.
fake._results.completed.add("h1")
s._sessions[("h", 1)] = fake
# Step 2: poll drains completion -> entry marked ready, request admitted.
assert s.ensure_cache_available(req, 0) is True
assert "h1" not in s._in_flight
assert "h1" in s._step_completed
entry = s._cache.get("h1")
assert entry is not None and entry.ready
# The scheduler now allocates encoder input; the ready entry loads through
# the same local path as a natively cached encoding.
s.update_state_after_alloc(req, 0)
meta = s.build_connector_meta(scheduler_output=None)
assert "h1" in meta.loads
assert meta.loads["h1"][2] == (1, width)
assert "h1" not in s._step_completed # cleared by promote
# Step 3: still cached -> admitted directly, no new transfer.
assert s.ensure_cache_available(req, 0) is True
s.shutdown()
@pytest.mark.parametrize("params", [None, {"x": 123}])
def test_no_remote_announcement_is_ready(monkeypatch, params):
s = _consumer_sched(monkeypatch)
assert s.ensure_cache_available(_Request([_Feature("x")], params=params), 0)
assert s.take_unavailable_requests() == set()
s.shutdown()
@pytest.mark.parametrize("payload", [None, "metadata", "pixels", "embeds", "shm"])
def test_failed_read_falls_back_only_with_local_input(monkeypatch, payload):
"""A remote miss is fatal only when local model input is unavailable."""
s = _consumer_sched(monkeypatch)
fake = _FakeSession()
def _fake_start(mm_hash, info, size, shape):
entry = s._cache.alloc(mm_hash, 1, shape)
assert entry is not None
fake.started.append(mm_hash)
return True
monkeypatch.setattr(s, "_start_xfer", _fake_start)
fields = {
"metadata": {"image_grid_thw": torch.tensor([1, 2, 2])},
"pixels": {"pixel_values": torch.ones(1, 3)},
"embeds": {"image_embeds": torch.ones(1, _HID)},
"shm": {"address": 123, "monotonic_id": 1},
}
data = (
MultiModalKwargsItem(
{
key: MultiModalFieldElem(value, MultiModalSharedField(batch_size=1))
for key, value in fields[payload].items()
}
)
if payload is not None
else None
)
req = _Request([_Feature("h1", 1, data)], params=_params("h1", 1))
# Step 1: start the read.
assert s.ensure_cache_available(req, 0) is False
s.build_connector_meta(scheduler_output=None)
fake._results.tombstoned.add("h1")
s._sessions[("h", 1)] = fake
can_fallback = payload in ("pixels", "embeds", "shm")
assert s.ensure_cache_available(req, 0) is can_fallback
assert s._cache.get("h1") is None
assert "h1" not in s._in_flight
assert s.take_unavailable_requests() == (set() if can_fallback else {"r1"})
# Drained by the read, so the scheduler cannot abort it twice.
assert s.take_unavailable_requests() == set()
if can_fallback:
# A later scheduling step must not retry the failed remote source.
s.build_connector_meta(scheduler_output=None)
assert s.ensure_cache_available(req, 0)
assert fake.started == ["h1"]
s.shutdown()
@pytest.mark.parametrize("retry", [False, True], ids=["in-flight", "retryable"])
def test_remote_wait_budget_survives_retries_and_long_steps(monkeypatch, retry):
"""Neither retries nor delayed scheduling may restart the request budget."""
s = _consumer_sched(monkeypatch)
fake = _FakeSession()
s._sessions[("h", 1)] = fake
def start(mm_hash, info, size, shape):
assert s._cache.alloc(mm_hash, 1, shape) is not None
fake.started.append(mm_hash)
return True
monkeypatch.setattr(s, "_start_xfer", start)
now = [1000.0]
monkeypatch.setattr(time, "monotonic", lambda: now[0])
req = _Request([_Feature("h1")], params=_params("h1", 1))
assert not s.ensure_cache_available(req, 0)
for elapsed in (20, 40, 121):
now[0] = 1000.0 + elapsed
if retry:
fake._results.retryable.add("h1")
s.build_connector_meta(scheduler_output=None)
assert not s.ensure_cache_available(req, 0)
assert s.take_unavailable_requests() == ({"r1"} if elapsed > 60 else set())
assert len(fake.started) == (3 if retry else 1)
s.shutdown()
def test_retryable_read_re_requests_without_admitting(monkeypatch):
"""A not-ready NACK defers the request and re-requests the read.
Unlike a tombstone it never admits the request, because an admitted request
whose media was rewritten away upstream has no embedding and no way to
recompute one.
"""
s = _consumer_sched(monkeypatch)
fake = _FakeSession()
def _fake_start(mm_hash, info, size, shape):
entry = s._cache.alloc(mm_hash, 1, shape)
assert entry is not None
fake.started.append(mm_hash)
return True
monkeypatch.setattr(s, "_start_xfer", _fake_start)
req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
assert s.ensure_cache_available(req, 0) is False
s.build_connector_meta(scheduler_output=None)
# Producer's save has not landed: retryable, not tombstoned.
fake._results.retryable.add("h1")
s._sessions[("h", 1)] = fake
# The entry is released and the read re-issued in the same step, so the
# request stays deferred rather than being admitted unfetched.
assert s.ensure_cache_available(req, 0) is False
assert "h1" in s._in_flight
assert "h1" not in s._tombstones
assert fake.started.count("h1") == 2
s.shutdown()
@pytest.mark.parametrize(
"fields",
[
{"size_bytes": 999},
{"size_bytes": None},
{"size_bytes": "big"},
{"size_bytes": _MISSING},
{"shape": _MISSING},
{"shape": [2, _HID]},
],
ids=[
"size-mismatch",
"null-size",
"text-size",
"no-size",
"no-shape",
"wrong-rows",
],
)
def test_invalid_announcement_fails_the_request(monkeypatch, fields):
"""An unusable announced size or shape fails the request, not the engine."""
s = _consumer_sched(monkeypatch)
announced = {**_params("h1", 1)["h1"], **fields}
announced = {k: v for k, v in announced.items() if v is not _MISSING}
req = _Request([_Feature("h1", 1)], params={"h1": announced})
assert s.ensure_cache_available(req, 0) is False
assert s.take_unavailable_requests() == {"r1"}
assert "h1" not in s._in_flight
assert s._cache.get("h1") is None
s.shutdown()
def test_deferral_budget_is_per_request(monkeypatch):
"""Shared hashes have independent wait budgets and cancellation cleanup."""
s = _consumer_sched(monkeypatch)
monkeypatch.setattr(s._cache, "alloc", lambda key, n, shape=None: None)
params = _params("h1", 1)
old = _Request([_Feature("h1", 1)], params=params, req_id="old")
new = _Request([_Feature("h1", 1)], params=params, req_id="new")
now = [1000.0]
monkeypatch.setattr(time, "monotonic", lambda: now[0])
assert s.ensure_cache_available(old, 0) is False
now[0] += sched_mod._ADMIT_DEFER_TIMEOUT_S - 1
assert s.ensure_cache_available(new, 0) is False
assert s.take_unavailable_requests() == set()
now[0] += 2
# `old` is now past its budget; `new` is not.
assert s.ensure_cache_available(new, 0) is False
assert s.take_unavailable_requests() == set()
assert s.request_finished(new) == (False, None)
assert set(s._deferred_since) == {("old", "h1")}
assert s.ensure_cache_available(old, 0) is False
assert s.take_unavailable_requests() == {"old"}
assert "h1" not in s._in_flight
s.shutdown()
def test_metadata_only_entry_admits_without_transfer(monkeypatch):
# A producer that reported placeholder metadata but had no cache entry
# to transfer omits peer_host/peer_port/size_bytes entirely; the consumer
# must fall back to local compute rather than treating that as a size
# mismatch.
s = _consumer_sched(monkeypatch)
metadata_only = {"h1": {"image_grid_thw": [1, 2, 3]}}
req = _Request([_Feature("h1", 1)], params=metadata_only)
assert s.ensure_cache_available(req, 0) is True
assert "h1" not in s._in_flight
assert s._cache.get("h1") is None
s.shutdown()
def test_orphan_not_ready_entry_defers_no_realloc(monkeypatch):
# Post-quarantine, post-tombstone-consumed orphan: a not-ready cache
# entry exists for the mm_hash, but the hash is tracked in none of
# _in_flight / _step_completed / _tombstones. Its blocks are held by a
# quarantined/settling DMA and must not be re-allocated.
s = _consumer_sched(monkeypatch)
entry = s._cache.alloc("h1", 1)
assert entry is not None and not entry.ready
assert "h1" not in s._in_flight
assert "h1" not in s._step_completed
assert "h1" not in s._tombstones
# Spy on alloc: record calls, delegate to the real allocator (which
# asserts on a duplicate key). A re-alloc for "h1" is the bug.
real_alloc = s._cache.alloc
calls: list[str] = []
def _spy_alloc(key, n):
calls.append(key)
return real_alloc(key, n)
monkeypatch.setattr(s._cache, "alloc", _spy_alloc)
req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
# Must not raise; the request is deferred while the DMA settles, the
# orphan entry is untouched, and alloc is never called again.
assert s.ensure_cache_available(req, 0) is False
assert calls == []
assert s._cache.get("h1") is entry
assert not entry.ready
assert s.take_unavailable_requests() == set()
s.shutdown()
def test_already_computed_feature_admits_without_transfer(monkeypatch):
# A feature entirely within num_computed_tokens is skipped by admit; no
# cache lookup/alloc or transfer should be triggered for it.
s = _consumer_sched(monkeypatch)
calls: list[str] = []
real_alloc = s._cache.alloc
def _spy_alloc(key, n):
calls.append(key)
return real_alloc(key, n)
monkeypatch.setattr(s._cache, "alloc", _spy_alloc)
# Feature spans [0, 4); with num_computed_tokens=4 it is fully computed.
req = _Request([_Feature("h1", 4)], params=_params("h1", 4))
assert s.ensure_cache_available(req, 4) is True
assert "h1" not in s._in_flight
assert calls == []
s.shutdown()
def test_local_ready_entry_admits_without_transfer(monkeypatch):
# A READY entry already present in the cache is a local hit: admit
# returns True and no new transfer is started for that mm_hash.
s = _consumer_sched(monkeypatch)
entry = s._cache.alloc("h1", 1)
assert entry is not None
s._cache.mark_ready("h1")
req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
assert s.ensure_cache_available(req, 0) is True
assert "h1" not in s._in_flight
s.shutdown()
def test_in_flight_hash_defers_without_second_transfer(monkeypatch):
# A mm_hash already tracked in _in_flight must defer the request rather
# than starting a duplicate transfer.
s = _consumer_sched(monkeypatch)
started: list[str] = []
def _spy_start_xfer(mm_hash, info, size, shape):
started.append(mm_hash)
return True
monkeypatch.setattr(s, "_start_xfer", _spy_start_xfer)
s._in_flight.add("h1")
req = _Request([_Feature("h1", 1)], params=_params("h1", 1))
assert s.ensure_cache_available(req, 0) is False
assert started == []
s.shutdown()
tests/v1/ec_connector/unit/test_scheduler_nixl_ctor.py
-146
View File
@@ -1,146 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import uuid
from pathlib import Path
import pytest
import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
import vllm.distributed.nixl_utils as nixl_utils
from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
from vllm.config.ec_transfer import ECTransferConfig
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
ECSharedRegion,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
from vllm.distributed.nixl_utils import NixlWrapper
_N = 16
_BS = 64
def _region() -> ECSharedRegion:
return ECSharedRegion(
engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS
)
def test_ec_enable_nixl_absent_from_extra_config_by_default():
cfg = ECTransferConfig()
assert cfg.get_from_extra_config("ec_enable_nixl", False) is False
def test_ec_enable_nixl_read_from_extra_config():
cfg = ECTransferConfig(
ec_connector="ECCPUConnector",
ec_role="ec_both",
ec_connector_extra_config={"ec_enable_nixl": True},
)
assert cfg.get_from_extra_config("ec_enable_nixl", False) is True
def test_gate_off_builds_no_nixl(monkeypatch):
monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region())
# ec_enable_nixl defaults to False.
s = ECCPUScheduler(create_ec_vllm_config(ec_role="ec_both"))
assert s._nixl_enabled is False
assert getattr(s, "_data", None) is None
assert getattr(s, "_producer_session", None) is None
s.shutdown()
def test_string_false_in_extra_config_leaves_nixl_off(monkeypatch):
"""Extra config is not type coerced, and bool("false") is True."""
monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region())
cfg = create_ec_vllm_config(ec_role="ec_both")
cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = "false"
s = ECCPUScheduler(cfg)
assert s._nixl_enabled is False
s.shutdown()
@pytest.mark.cpu_test
def test_missing_nixl_releases_shared_region(monkeypatch):
"""A failed constructor leaves no scheduler for the caller to shut down."""
# Import before patching to avoid caching the missing-NIXL values.
from vllm.distributed.ec_transfer.ec_connector.cpu.data import nixl # noqa: F401
region = _region()
monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)
monkeypatch.setattr(nixl_utils, "NixlWrapper", None)
monkeypatch.setattr(nixl_utils, "nixl_agent_config", None)
cfg = create_ec_vllm_config(ec_role="ec_both")
cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True
try:
with pytest.raises(RuntimeError, match="ec_enable_nixl requires NIXL"):
ECCPUScheduler(cfg)
assert not Path(region._mmap_path).exists()
assert region._fd is None
assert region._mmap_obj is None
finally:
region.cleanup()
@pytest.mark.cpu_test
def test_control_transport_failure_releases_nixl_and_shared_region(monkeypatch):
import vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq as zmq_mod
import vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl as data_mod
class RegisteredData:
registered = True
def deregister(self):
self.registered = False
data = RegisteredData()
error = RuntimeError("control transport failed")
def fail_control_transport():
raise error
region = _region()
monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: region)
monkeypatch.setattr(nixl_utils, "NixlWrapper", object())
monkeypatch.setattr(nixl_utils, "nixl_agent_config", object())
monkeypatch.setattr(data_mod, "NixlDataTransport", lambda **kwargs: data)
monkeypatch.setattr(zmq_mod, "ZmqClientTransport", fail_control_transport)
cfg = create_ec_vllm_config(ec_role="ec_consumer")
cfg.model_config.hf_config = None
cfg.model_config.get_inputs_embeds_size.return_value = 8
cfg.model_config.model = "test-model"
cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True
try:
with pytest.raises(RuntimeError, match="control transport failed") as exc:
ECCPUScheduler(cfg)
assert exc.value is error
assert not data.registered
assert not Path(region._mmap_path).exists()
assert region._fd is None
assert region._mmap_obj is None
finally:
region.cleanup()
@pytest.mark.skipif(NixlWrapper is None, reason="Requires NIXL package")
def test_gate_on_wires_data_transport_and_producer_session(monkeypatch):
# Port 0 lets the OS pick an ephemeral port so the real ZMQ ROUTER bind
# in ProducerSession.start() cannot collide with another test/process.
monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_HOST", "127.0.0.1")
monkeypatch.setenv("VLLM_EC_SIDE_CHANNEL_PORT", "0")
monkeypatch.setattr(sched_mod, "create_ec_shared_region", lambda cfg: _region())
cfg = create_ec_vllm_config(ec_role="ec_both")
cfg.ec_transfer_config.ec_connector_extra_config["ec_enable_nixl"] = True
s = ECCPUScheduler(cfg)
try:
assert s._nixl_enabled is True
assert s._data is not None
assert s._compat_hash is not None
assert s._producer_session is not None
assert s._peer_host == "127.0.0.1"
assert s._transport is not None # is_ec_consumer is also True
finally:
s.shutdown()
tests/v1/ec_connector/unit/test_scheduler_nixl_producer.py
-116
View File
@@ -1,116 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import uuid
import vllm.distributed.ec_transfer.ec_connector.cpu.scheduler as sched_mod
from tests.v1.ec_connector.unit.utils import create_ec_vllm_config
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
ECSharedRegion,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import ECCPUScheduler
_N, _BS = 16, 64
class _Pos:
def __init__(self, offset, length):
self.offset, self.length = offset, length
def get_num_embeds(self):
return self.length
class _Feature:
def __init__(self, mm_hash, length=1):
self.mm_hash = mm_hash
self.identifier = mm_hash
self.mm_position = _Pos(0, length)
self.data = None
self.modality = "image"
class _Request:
def __init__(self, features, req_id="r1"):
self.mm_features = features
self.request_id = req_id
def _sched_gate_off(monkeypatch):
def _region(cfg):
return ECSharedRegion(
engine_id="eng-" + str(uuid.uuid4()), num_blocks=_N, block_size_bytes=_BS
)
monkeypatch.setattr(sched_mod, "create_ec_shared_region", _region)
return ECCPUScheduler(create_ec_vllm_config(ec_role="ec_both"))
def test_request_finished_gate_off_returns_none(monkeypatch):
s = _sched_gate_off(monkeypatch)
assert s.request_finished(_Request([_Feature("h1")])) == (False, None)
s.shutdown()
def test_request_finished_producer_emits_params(monkeypatch):
s = _sched_gate_off(monkeypatch)
# Simulate NIXL-enabled producer bookkeeping without building real NIXL.
s._nixl_enabled = True
s._peer_host, s._peer_port = "1.2.3.4", 5601
# feature length=2, hidden_dim=32, element_size=2 -> 128 bytes -> 2 blocks.
entry = s._cache.alloc("h1", 2, (2, 32))
assert entry is not None
s._cache.mark_ready("h1")
delay, params = s.request_finished(_Request([_Feature("h1", length=2)]))
assert delay is False
assert entry.evictable
assert params == {
"h1": {
"metadata": {},
"item_indices": [0],
"peer_host": "1.2.3.4",
"peer_port": 5601,
"size_bytes": 2 * 32 * 2,
"shape": [2, 32],
}
}
s.shutdown()
def test_request_finished_announces_not_ready_entry(monkeypatch):
"""A save whose GPU->mmap copy hasn't been confirmed complete yet is
still announced: the entry can't be evicted before it's ready, and a read
arriving too early is NACKed NACK_NOT_READY for the consumer to retry —
protecting it is not this method's job."""
s = _sched_gate_off(monkeypatch)
s._nixl_enabled = True
s._peer_host, s._peer_port = "1.2.3.4", 5601
s._cache.alloc("h1", 2, (2, 32)) # allocated but not marked ready
delay, params = s.request_finished(_Request([_Feature("h1", length=2)]))
assert delay is False
assert params == {
"h1": {
"metadata": {},
"item_indices": [0],
"peer_host": "1.2.3.4",
"peer_port": 5601,
"size_bytes": 2 * 32 * 2,
"shape": [2, 32],
}
}
s.shutdown()
def test_request_finished_skips_unallocated_entry(monkeypatch):
s = _sched_gate_off(monkeypatch)
s._nixl_enabled = True
s._peer_host, s._peer_port = "1.2.3.4", 5601
# No alloc() for "h1" at all — e.g. the cache was full at save time.
delay, params = s.request_finished(_Request([_Feature("h1", length=2)]))
assert delay is False
# The item's placeholder metadata is still reported even though there's
# no cache entry to transfer (empty "metadata": no fields, no transfer).
assert params == {"h1": {"metadata": {}, "item_indices": [0]}}
s.shutdown()
tests/v1/ec_connector/unit/test_session.py
-471
View File
@@ -1,471 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for ProducerXfer, ProducerSession, ConsumerXfer, ConsumerSession."""
import time
from unittest.mock import MagicMock
import pytest
from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import (
XferAck,
XferReq,
XferStatus,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
EmbeddingCache,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.session import (
_CONSUMER_QUARANTINE_TIMEOUT_S,
ConsumerSession,
ConsumerXfer,
ProducerSession,
ProducerXfer,
XferState,
)
# ── helpers ───────────────────────────────────────────────────────────────────
def _make_data(xfer_state: str = "PROC") -> MagicMock:
data = MagicMock()
data.check_xfer_state.return_value = xfer_state
data.post_read.return_value = 99
return data
def _ok_ack(
mm_hash: str = "h1",
session_id: str = "sess-1",
src_indices: list | None = None,
) -> XferAck:
return XferAck(
mm_hash=mm_hash,
status=XferStatus.OK,
session_id=session_id,
src_block_indices=src_indices or [0, 1],
agent_metadata=b"meta",
mem_descriptor=b"desc",
)
def _xfer(
data: MagicMock | None = None,
deadline: float | None = None,
consumer_session_id: str = "sess-1",
) -> ConsumerXfer:
return ConsumerXfer(
mm_hash="h1",
block_indices=[10, 11],
addr=("host", 1234),
deadline=deadline if deadline is not None else time.monotonic() + 60,
data=data or _make_data(),
consumer_session_id=consumer_session_id,
)
def _started(
data: MagicMock | None = None, deadline: float | None = None
) -> ConsumerXfer:
"""ConsumerXfer that has already received an OK XferAck."""
x = _xfer(data, deadline)
x.handle_ack(_ok_ack(session_id="sess-1", src_indices=[5, 6]), agent_name="agent-1")
return x
# ── ProducerXfer ──────────────────────────────────────────────────────────────
def test_producer_xfer_not_expired_before_deadline():
x = ProducerXfer("h1", [0, 1], deadline=time.monotonic() + 60)
assert not x.is_expired()
def test_producer_xfer_expired_after_deadline():
x = ProducerXfer("h1", [0, 1], deadline=time.monotonic() - 1)
assert x.is_expired()
# ── ProducerSession ───────────────────────────────────────────────────────────
def _xfer_req(mm_hash: str = "h1", session_id: str = "sess-1") -> XferReq:
return XferReq(
mm_hash=mm_hash, compatibility_hash="hash-abc", session_id=session_id
)
def _make_producer_session(cache: EmbeddingCache | None = None) -> ProducerSession:
data = _make_data()
data.get_agent_metadata.return_value = b"meta"
data.get_mem_descriptor.return_value = b"desc"
return ProducerSession(
transport=MagicMock(),
data=data,
cache=cache or EmbeddingCache(num_blocks=8),
compat_hash="hash-abc",
)
def test_producer_session_grant_pins_ready_entry_and_returns_blocks():
cache = EmbeddingCache(num_blocks=8)
cache.alloc("h1", 2)
cache.mark_ready("h1")
s = _make_producer_session(cache)
ack = s._grant_or_nack(_xfer_req(mm_hash="h1"))
assert ack.status == XferStatus.OK
entry = cache.get("h1")
assert not entry.evictable # pinned by the grant
assert list(ack.src_block_indices) == list(entry.block_ids)
def test_producer_session_grant_nacks_missing_entry():
s = _make_producer_session()
ack = s._grant_or_nack(_xfer_req(mm_hash="nope"))
assert ack.status == XferStatus.NACK_MISSING
def test_producer_session_grant_nacks_not_ready_entry_as_retryable():
"""An announced entry whose save is still in flight is distinguished from a
miss, so the consumer retries instead of falling back to local compute."""
cache = EmbeddingCache(num_blocks=8)
cache.alloc("h1", 2) # not ready
s = _make_producer_session(cache)
ack = s._grant_or_nack(_xfer_req(mm_hash="h1"))
assert ack.status == XferStatus.NACK_NOT_READY
def test_producer_session_grant_nacks_version_mismatch():
cache = EmbeddingCache(num_blocks=8)
cache.alloc("h1", 2)
cache.mark_ready("h1")
s = _make_producer_session(cache)
req = XferReq(
mm_hash="h1",
compatibility_hash="hash-abc",
session_id="sess-1",
connector_version=-1,
)
ack = s._grant_or_nack(req)
assert ack.status == XferStatus.NACK_VERSION
entry = cache.get("h1")
assert entry.evictable # nothing was pinned
def test_producer_session_grant_nacks_compat_hash_mismatch():
cache = EmbeddingCache(num_blocks=8)
cache.alloc("h1", 2)
cache.mark_ready("h1")
s = _make_producer_session(cache)
req = XferReq(mm_hash="h1", compatibility_hash="wrong-hash", session_id="sess-1")
ack = s._grant_or_nack(req)
assert ack.status == XferStatus.NACK_INCOMPAT
entry = cache.get("h1")
assert entry.evictable # nothing was pinned
def test_producer_session_notif_unpins_via_cache():
cache = EmbeddingCache(num_blocks=8)
cache.alloc("h1", 2)
cache.mark_ready("h1")
s = _make_producer_session(cache)
s._grant_or_nack(_xfer_req(mm_hash="h1", session_id="sess-1"))
entry = cache.get("h1")
assert not entry.evictable # pinned
s._data.get_new_notifs.return_value = {"agent": [b"sess-1:h1"]}
s._drain_notifs()
assert entry.evictable # unpinned via cache
assert s._active_xfers == {}
def test_producer_session_sweep_timeouts_unpins_via_cache():
cache = EmbeddingCache(num_blocks=8)
cache.alloc("h1", 2)
cache.mark_ready("h1")
s = _make_producer_session(cache)
s._grant_or_nack(_xfer_req(mm_hash="h1", session_id="sess-1"))
key = next(iter(s._active_xfers))
s._active_xfers[key].deadline = time.monotonic() - 1
s._sweep_timeouts()
assert cache.get("h1").evictable
assert s._active_xfers == {}
# ── ConsumerXfer.handle_ack ───────────────────────────────────────────────────
def test_consumer_xfer_handle_ack_ok_starts_read():
data = _make_data()
x = _xfer(data)
assert x.handle_ack(_ok_ack(), agent_name="agent-1") is True
assert x.transfer_handle == 99
def test_consumer_xfer_handle_ack_notif_msg_is_session_id_colon_mm_hash():
data = _make_data()
x = _xfer(data, consumer_session_id="my-sess")
x.handle_ack(_ok_ack(mm_hash="h1"), agent_name="agent-1")
_, kwargs = data.post_read.call_args
assert kwargs["notif_msg"] == b"my-sess:h1"
def test_consumer_xfer_handle_ack_forwards_correct_args():
data = _make_data()
x = _xfer(data)
x.handle_ack(_ok_ack(src_indices=[5, 6]), agent_name="agent-1")
args, _ = data.post_read.call_args
assert args[0] == [10, 11] # local block_indices
assert args[1] == "agent-1" # agent_name (not a raw handle)
assert args[2] == [5, 6] # remote src_block_indices
def test_consumer_xfer_handle_ack_resets_deadline():
original = time.monotonic() + 2
x = _xfer(deadline=original)
x.handle_ack(_ok_ack(), agent_name="agent-1")
assert x.deadline > original
def test_consumer_xfer_handle_ack_nack_returns_false():
data = _make_data()
x = _xfer(data)
assert (
x.handle_ack(XferAck(mm_hash="h1", status=XferStatus.NACK_MISSING), "a")
is False
)
assert x.transfer_handle is None
data.post_read.assert_not_called()
# ── ConsumerXfer.poll — WAITING_ACK ──────────────────────────────────────────
def test_consumer_xfer_poll_waiting_ack_before_deadline():
assert _xfer().poll(time.monotonic()) == XferState.WAITING_ACK
def test_consumer_xfer_poll_waiting_ack_timeout():
assert (
_xfer(deadline=time.monotonic() - 1).poll(time.monotonic())
== XferState.ACK_TIMEOUT
)
# ── ConsumerXfer.poll — READING ───────────────────────────────────────────────
def test_consumer_xfer_poll_reading_proc():
assert _started().poll(time.monotonic()) == XferState.READING
def test_consumer_xfer_poll_reading_done_releases_handle():
data = _make_data(xfer_state="DONE")
x = _started(data)
assert x.poll(time.monotonic()) == XferState.DONE
data.release_xfer_handle.assert_called_once_with(99)
assert x.transfer_handle is None
def test_consumer_xfer_poll_reading_unexpected_state_fails():
data = _make_data(xfer_state="ERR")
x = _started(data)
assert x.poll(time.monotonic()) == XferState.READ_FAILED
data.release_xfer_handle.assert_called_once_with(99)
def test_consumer_xfer_poll_reading_exception_fails_gracefully():
data = _make_data()
data.check_xfer_state.side_effect = RuntimeError("nixl crash")
x = _started(data)
assert x.poll(time.monotonic()) == XferState.READ_FAILED
def test_consumer_xfer_poll_timeout_quarantines_without_releasing():
data = _make_data(xfer_state="PROC")
x = _started(data)
assert x.poll(time.monotonic() + 9999) == XferState.QUARANTINED
data.release_xfer_handle.assert_not_called()
assert x.transfer_handle == 99
# ── ConsumerXfer.poll — QUARANTINED ───────────────────────────────────────────
def test_consumer_xfer_quarantined_proc_stays():
data = _make_data(xfer_state="PROC")
x = _started(data)
x.poll(time.monotonic() + 9999) # → QUARANTINED
assert x.poll(time.monotonic()) == XferState.QUARANTINED
data.release_xfer_handle.assert_not_called()
def test_consumer_xfer_quarantined_terminal_settles():
data = _make_data(xfer_state="PROC")
x = _started(data)
x.poll(time.monotonic() + 9999) # → QUARANTINED
data.check_xfer_state.return_value = "DONE"
assert x.poll(time.monotonic()) == XferState.SETTLED
data.release_xfer_handle.assert_called_once_with(99)
assert x.transfer_handle is None
def test_consumer_xfer_quarantined_exception_settles():
data = _make_data(xfer_state="PROC")
x = _started(data)
x.poll(time.monotonic() + 9999)
data.check_xfer_state.side_effect = RuntimeError("gone")
assert x.poll(time.monotonic()) == XferState.SETTLED
def test_consumer_xfer_quarantined_timeout_evicts_blocks():
"""Quarantine timeout should release handle and return READ_FAILED."""
data = _make_data(xfer_state="PROC")
x = _started(data)
# First poll transitions to QUARANTINED (deadline extended by 60s)
x.poll(time.monotonic() + 9999)
# Second poll with expired quarantine deadline should evict
expired_time = time.monotonic() + 9999 + _CONSUMER_QUARANTINE_TIMEOUT_S + 1
assert x.poll(expired_time) == XferState.READ_FAILED
data.release_xfer_handle.assert_called_once_with(99)
assert x.transfer_handle is None
# ── ConsumerXfer.cancel / release ────────────────────────────────────────────
def test_consumer_xfer_cancel_valid_in_waiting_ack():
_xfer().cancel() # must not raise
def test_consumer_xfer_cancel_raises_when_read_started():
with pytest.raises(AssertionError):
_started().cancel()
def test_consumer_xfer_release_releases_handle():
data = _make_data()
x = _started(data)
x.release()
data.release_xfer_handle.assert_called_once_with(99)
def test_consumer_xfer_release_noop_when_no_handle():
data = _make_data()
_xfer(data).release()
data.release_xfer_handle.assert_not_called()
# ── ConsumerSession ───────────────────────────────────────────────────────────
def _make_consumer_session(data=None):
zmq_conn = MagicMock()
zmq_conn.recv.return_value = []
return ConsumerSession(
addr=("host", 1234),
zmq_conn=zmq_conn,
transport=MagicMock(),
data=data or _make_data(),
compat_hash="hash-abc",
)
def test_consumer_session_has_stable_session_id():
s = _make_consumer_session()
assert isinstance(s._session_id, str) and len(s._session_id) > 0
assert s._session_id == s._session_id # stable
def test_consumer_session_start_xfer_sends_req_with_session_id():
import msgspec
s = _make_consumer_session()
s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10)
raw = s._zmq.send.call_args[0][0]
req = msgspec.msgpack.decode(raw, type=XferReq)
assert req.mm_hash == "h1"
assert req.session_id == s._session_id
def test_consumer_session_poll_done_goes_to_completed():
import msgspec
data = _make_data(xfer_state="PROC")
s = _make_consumer_session(data)
s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10)
ack = XferAck(
mm_hash="h1",
status=XferStatus.OK,
session_id=s._session_id,
src_block_indices=[5, 6],
agent_metadata=b"meta",
mem_descriptor=b"desc",
)
data.add_remote_peer.return_value = "agent-1"
raw_ack = msgspec.msgpack.encode(ack)
data.check_xfer_state.return_value = "DONE"
s.poll([raw_ack], time.monotonic())
results = s.take_results()
assert "h1" in results.completed
assert not results.tombstoned
def test_consumer_session_poll_nack_goes_to_tombstoned():
import msgspec
s = _make_consumer_session()
s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10)
nack = XferAck(mm_hash="h1", status=XferStatus.NACK_MISSING)
s.poll([msgspec.msgpack.encode(nack)], time.monotonic())
results = s.take_results()
assert "h1" in results.tombstoned
assert not results.completed
def test_consumer_session_poll_not_ready_nack_goes_to_retryable():
"""A not-ready NACK must not tombstone: tombstoning admits the request with
no embedding, which is fatal once the media has been rewritten away."""
import msgspec
s = _make_consumer_session()
s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 10)
nack = XferAck(mm_hash="h1", status=XferStatus.NACK_NOT_READY)
s.poll([msgspec.msgpack.encode(nack)], time.monotonic())
results = s.take_results()
assert "h1" in results.retryable
assert not results.tombstoned
assert not results.completed
def test_consumer_session_ack_timeout_reports_retry_once():
"""A delayed producer ACK is retryable, and its result is drained once."""
s = _make_consumer_session()
s.start_xfer("h1", [0, 1], deadline=time.monotonic() - 1)
s.poll([], time.monotonic())
r1 = s.take_results()
assert "h1" in r1.retryable
assert not r1.tombstoned
r2 = s.take_results()
assert not r2.retryable # cleared after first take
def test_consumer_session_on_peer_down_cancels_waiting_ack():
s = _make_consumer_session()
s.start_xfer("h1", [0, 1], deadline=time.monotonic() + 60)
s.on_peer_down()
results = s.take_results()
assert "h1" in results.retryable
assert not results.tombstoned
tests/v1/ec_connector/unit/test_utils.py
-54
View File
@@ -1,54 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for ECCPUConnector scheduler utilities."""
import msgspec
import pytest
from vllm.distributed.ec_transfer.ec_connector.cpu.utils import (
build_block_descs,
deserialize_mem_descriptor,
serialize_mem_descriptor,
)
# ── build_block_descs ────────────────────────────────────────────────────────
@pytest.mark.parametrize("device_id,expected_dev", [(7, 7), (None, 0)])
def test_build_block_descs(device_id, expected_dev):
kwargs = dict(base_ptr=1000, num_blocks=4, block_size_bytes=256)
if device_id is not None:
kwargs["device_id"] = device_id
descs = build_block_descs(**kwargs)
assert len(descs) == 4
for i, (addr, size, dev) in enumerate(descs):
assert addr == 1000 + i * 256
assert size == 256
assert dev == expected_dev
def test_build_block_descs_zero_blocks_returns_empty():
assert build_block_descs(base_ptr=100, num_blocks=0, block_size_bytes=64) == []
# ── serialize / deserialize_mem_descriptor ───────────────────────────────────
def test_mem_descriptor_roundtrip():
descs = [(100, 64, 0), (164, 64, 0), (228, 64, 1)]
assert deserialize_mem_descriptor(serialize_mem_descriptor(descs)) == descs
@pytest.mark.parametrize(
"bad_value",
[
[(1, 2)], # 2-tuple instead of 3
[("a", "b", "c")], # strings, not ints
],
)
def test_mem_descriptor_rejects_malformed_payload(bad_value):
"""Malformed descriptor lists must fail to decode."""
encoder = msgspec.msgpack.Encoder()
bad_payload = encoder.encode(bad_value)
with pytest.raises((msgspec.DecodeError, msgspec.ValidationError)):
deserialize_mem_descriptor(bad_payload)
tests/v1/ec_connector/unit/test_worker_ec_connector.py
-115
View File
@@ -1,115 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for the V2 GPU model runner's EC connector wrapper."""
from types import SimpleNamespace
from unittest.mock import MagicMock, patch
import pytest
from vllm.distributed.ec_transfer.ec_connector.base import (
ECConnectorBase,
ECConnectorMetadata,
)
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
from vllm.v1.worker.ec_connector_model_runner_mixin import (
ECConnectorModelRunnerMixin,
)
from vllm.v1.worker.gpu.ec_connector import NO_OP_EC_CONNECTOR, ActiveECConnector
pytestmark = pytest.mark.cpu_test
WORKER_META = object()
CONNECTOR_STATS = object()
def _scheduler_output() -> SimpleNamespace:
return SimpleNamespace(
ec_connector_metadata=ECConnectorMetadata(), finished_req_ids=frozenset()
)
def _connector(
encoder_cache: dict | None = None,
is_producer: bool = True,
is_consumer: bool = False,
) -> tuple[ActiveECConnector, MagicMock]:
fake = MagicMock(spec=ECConnectorBase)
fake.is_producer = is_producer
fake.is_consumer = is_consumer
fake.get_finished.return_value = (None, None)
fake.build_connector_worker_meta.return_value = WORKER_META
fake.get_ec_connector_stats.return_value = CONNECTOR_STATS
with patch("vllm.v1.worker.gpu.ec_connector.get_ec_transfer", return_value=fake):
return ActiveECConnector(SimpleNamespace(), encoder_cache or {}), fake
@pytest.mark.parametrize(
("is_producer", "is_consumer"), [(True, False), (True, True), (False, True)]
)
def test_saves_newly_added_caches_for_every_producer(is_producer, is_consumer):
"""An ec_both node is also a producer: it must offload what it just computed."""
encoder_cache = {"mm_old": None}
connector, fake = _connector(encoder_cache, is_producer, is_consumer)
with connector.maybe_get_output(_scheduler_output()):
encoder_cache["mm_new"] = None
saved = [call.kwargs["mm_hash"] for call in fake.save_caches.call_args_list]
assert saved == (["mm_new"] if is_producer else [])
assert fake.start_save_caches.called == is_producer
if is_producer:
assert fake.start_save_caches.call_args.kwargs["encoder_cache"] is encoder_cache
assert fake.start_load_caches.called == is_consumer
def test_worker_meta_is_reported_on_context_exit():
"""Reported in the finally block, so is_empty() sees it only after the exit."""
connector, fake = _connector()
with connector.maybe_get_output(_scheduler_output()) as output:
assert output.ec_connector_worker_meta is None
assert output.ec_connector_worker_meta is WORKER_META
assert fake.clear_connector_metadata.called
def test_connector_stats_are_reported_on_context_exit():
"""Reported in the finally block, alongside the worker meta."""
connector, fake = _connector()
with connector.maybe_get_output(_scheduler_output()) as output:
assert output.ec_connector_stats is None
assert output.ec_connector_stats is CONNECTOR_STATS
def test_v1_producer_receives_existing_encoder_cache():
encoder_cache = {"mm_cached": None}
fake = MagicMock(spec=ECConnectorBase)
fake.is_producer = True
fake.is_consumer = False
fake.get_finished.return_value = (None, None)
module = "vllm.v1.worker.ec_connector_model_runner_mixin"
with (
patch(f"{module}.has_ec_transfer", return_value=True),
patch(f"{module}.get_ec_transfer", return_value=fake),
ECConnectorModelRunnerMixin.maybe_get_ec_connector_output(
_scheduler_output(), encoder_cache
),
):
pass
assert fake.start_save_caches.call_args.kwargs["encoder_cache"] is encoder_cache
def test_no_forward_reports_without_running_the_model():
connector, _ = _connector()
output = connector.no_forward(_scheduler_output())
assert output.ec_connector_output.ec_connector_worker_meta is WORKER_META
empty = NO_OP_EC_CONNECTOR.no_forward(_scheduler_output())
assert empty is EMPTY_MODEL_RUNNER_OUTPUT
tests/v1/ec_connector/unit/utils.py
-81
View File
@@ -1,81 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared config builders for EC connector unit tests."""
import uuid
from unittest.mock import Mock
import torch
from vllm.config import ModelConfig, VllmConfig
from vllm.config.ec_transfer import ECRole, ECTransferConfig
from vllm.config.parallel import ParallelConfig
# The scheduler derives its per-load ack count from the PCP size only under the
# `mp` backend -- which also keeps the world size off the local GPU count, so
# TP > 1 stays expressible on a single-GPU host.
_EXECUTOR_BACKEND = "mp"
_PCP_SIZE = 1
_EMBED_WIDTH = 32
def create_ec_vllm_config(
*,
ec_role: ECRole = "ec_both",
tensor_parallel_size: int = 1,
rank: int = 0,
dtype: torch.dtype = torch.float16,
encoder_output_widths: dict[str, int] | None = None,
) -> Mock:
"""Build a `VllmConfig` stand-in for EC connector unit tests.
`ec_transfer_config` and `parallel_config` are the real config objects, so
role derivation and the rank arithmetic the connector depends on behave as
they do in production.
`model_config` supplies a dtype and embedding width without resolving an
HF model or inspecting its architecture.
Args:
ec_role: EC role this instance plays. `is_ec_producer` and
`is_ec_consumer` are derived from it, as in production.
tensor_parallel_size: TP degree.
rank: Global rank of this worker.
dtype: Encoder cache dtype.
encoder_output_widths: Per-modality encoder output width, as a
producer measures it at startup. Defaults to the embedding width.
Returns:
A `VllmConfig`-specced mock carrying the configs above.
"""
parallel_config = ParallelConfig(
tensor_parallel_size=tensor_parallel_size,
prefill_context_parallel_size=_PCP_SIZE,
distributed_executor_backend=_EXECUTOR_BACKEND,
rank=rank,
)
ec_transfer_config = ECTransferConfig(
# The role properties report False unless a connector is named.
ec_connector="ECCPUConnector",
ec_role=ec_role,
engine_id=str(uuid.uuid4()),
)
ec_transfer_config.mm_encoder_output_widths = (
dict.fromkeys(("image", "video", "audio"), _EMBED_WIDTH)
if encoder_output_widths is None
else encoder_output_widths
)
model_config = Mock(spec=ModelConfig)
model_config.dtype = dtype
model_config.get_inputs_embeds_size.return_value = _EMBED_WIDTH
vllm_config = Mock(spec=VllmConfig)
vllm_config.ec_transfer_config = ec_transfer_config
vllm_config.parallel_config = parallel_config
vllm_config.model_config = model_config
vllm_config.instance_id = f"ec-test-{uuid.uuid4()}"
return vllm_config
vllm/config/kernel.py
+1 -2
View File
@@ -124,7 +124,6 @@ MoEBackend = Literal[
"triton",
"batched_triton",
"deep_gemm",
"deep_gemm_mega_moe",
"cutlass",
"flashinfer_trtllm",
"flashinfer_cutlass",
@@ -163,7 +162,7 @@ PASSTHROUGH_ALL2ALL_BACKEND: Final = "passthrough"
# "Native" here meaning that the backend does not use the fused MoE path.
# Currently only relevant for the vLLM provided DeepGEMM MegaMoE.,
# TODO: migrate to the fused MoE path ultimately.
NATIVE_MEGA_MOE_BACKENDS = frozenset({"deep_gemm_mega_moe"})
NATIVE_MEGA_MOE_BACKENDS: frozenset[str] = frozenset()
# Every megakernel backend. The FlashInfer variants run through the standard
# fused-MoE model path; the quantization method binds the kernel.
vllm/distributed/device_communicators/cpu_communicator.py
-339
View File
@@ -1,339 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from typing import Any
import torch
from torch.distributed import ProcessGroup
from vllm.distributed.utils import pickle
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.platforms.interface import CpuArchEnum
from .base_device_communicator import DeviceCommunicatorBase
logger = init_logger(__name__)
class CpuCommunicator(DeviceCommunicatorBase):
def __init__(
self,
cpu_group: ProcessGroup,
device: torch.device | None = None,
device_group: ProcessGroup | None = None,
unique_name: str = "",
use_all2all: bool = False,
):
super().__init__(
cpu_group, device, device_group, unique_name, use_all2all=use_all2all
)
self.dist_module = torch.distributed
if (
(
current_platform.get_cpu_architecture() == CpuArchEnum.X86
or current_platform.get_cpu_architecture() == CpuArchEnum.ARM
or current_platform.get_cpu_architecture() == CpuArchEnum.POWERPC
or current_platform.get_cpu_architecture() == CpuArchEnum.S390X
)
and hasattr(torch.ops._C, "init_shm_manager")
and (unique_name.startswith("tp") or unique_name.startswith("pp"))
and self._all_group_ranks_share_shm_group_name()
):
self.dist_module = _CPUSHMDistributed(self)
elif unique_name.startswith("tp") or unique_name.startswith("pp"):
logger.info(
"CPU SHM communicator disabled for group %s: ranks do not share "
"the same SHM group name, falling back to torch.distributed.",
unique_name,
)
# send/recv tensor_dict is only supported through the SHM communicator backend
self.supports_tensor_dict = isinstance(self.dist_module, _CPUSHMDistributed)
if self.use_all2all:
if self.all2all_backend not in (
"naive",
"allgather_reducescatter",
): # type: ignore[has-type]
logger.warning(
"`%s` all2all manager is not supported on CPU. "
"Falling back to `allgather_reducescatter` manager.",
self.all2all_backend, # type: ignore[has-type]
)
from .all2all import AgRsAll2AllManager
self.all2all_manager = AgRsAll2AllManager(self.cpu_group)
logger.info("Using allgather_reducescatter all2all manager.")
def _all_group_ranks_share_shm_group_name(self) -> bool:
"""CPUSHM requires all ranks in this group to agree on one SHM group name.
This is a lightweight consistency check for VLLM_DIST_IDENT/name inputs.
"""
local_name = _CPUSHMDistributed.make_group_name(self)
names: list[str] = [""] * self.world_size
torch.distributed.all_gather_object(
names,
local_name,
group=self.device_group,
)
return len(set(names)) == 1
def all_reduce(self, input_):
self.dist_module.all_reduce(input_, group=self.device_group)
return input_
def gather(
self, input_: torch.Tensor, dst: int = 0, dim: int = -1
) -> torch.Tensor | None:
"""NOTE: We assume that the input tensor is on the same device across
all the ranks.
NOTE: `dst` is the local rank of the destination rank.
"""
world_size = self.world_size
assert -input_.dim() <= dim < input_.dim(), (
f"Invalid dim ({dim}) for input tensor with shape {input_.size()}"
)
if dim < 0:
# Convert negative dim to positive.
dim += input_.dim()
# Allocate output tensor.
if self.rank_in_group == dst:
gather_list = [torch.empty_like(input_) for _ in range(world_size)]
else:
gather_list = None
# Gather.
self.dist_module.gather(
input_, gather_list, dst=self.ranks[dst], group=self.device_group
)
if self.rank_in_group == dst:
output_tensor = torch.cat(gather_list, dim=dim)
else:
output_tensor = None
return output_tensor
def all_gather(self, input_: torch.Tensor, dim: int = -1) -> torch.Tensor:
if dim < 0:
# Convert negative dim to positive.
dim += input_.dim()
input_size = input_.size()
# NOTE: we have to use concat-style all-gather here,
# stack-style all-gather has compatibility issues with
# torch.compile . see https://github.com/pytorch/pytorch/issues/138795
output_size = (input_size[0] * self.world_size,) + input_size[1:]
# Allocate output tensor.
output_tensor = torch.empty(
output_size, dtype=input_.dtype, device=input_.device
)
# All-gather.
self.dist_module.all_gather_into_tensor(
output_tensor, input_, group=self.device_group
)
# Reshape
output_tensor = output_tensor.reshape((self.world_size,) + input_size)
output_tensor = output_tensor.movedim(0, dim)
output_tensor = output_tensor.reshape(
input_size[:dim]
+ (self.world_size * input_size[dim],)
+ input_size[dim + 1 :]
)
return output_tensor
def send_tensor_dict(
self,
tensor_dict: dict[str, torch.Tensor | Any],
dst: int,
) -> None:
if not self.supports_tensor_dict:
raise NotImplementedError(
"CpuCommunicator does not support tensor dict fastpath with "
"torch.distributed backend."
)
return self.dist_module.send_tensor_dict(tensor_dict, dst)
def recv_tensor_dict(
self,
src: int,
) -> dict[str, torch.Tensor | Any]:
if not self.supports_tensor_dict:
raise NotImplementedError(
"CpuCommunicator does not support tensor dict fastpath with "
"torch.distributed backend."
)
return self.dist_module.recv_tensor_dict(src)
def dispatch_router_logits(
self,
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
is_sequence_parallel: bool = False,
extra_tensors: list[torch.Tensor] | None = None,
) -> (
tuple[torch.Tensor, torch.Tensor]
| tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]
):
"""Dispatch the hidden states and router logits to the appropriate device.
This is a no-op in the base class.
"""
assert self.all2all_manager is not None
return self.all2all_manager.dispatch_router_logits(
hidden_states,
router_logits,
is_sequence_parallel,
extra_tensors,
)
def dispatch(
self,
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
is_sequence_parallel: bool = False,
extra_tensors: list[torch.Tensor] | None = None,
) -> (
tuple[torch.Tensor, torch.Tensor, torch.Tensor]
| tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
):
"""Dispatch the hidden states and topk weights/ids to the appropriate device.
This is a no-op in the base class.
"""
assert self.all2all_manager is not None
return self.all2all_manager.dispatch(
hidden_states,
topk_weights,
topk_ids,
is_sequence_parallel,
extra_tensors=extra_tensors,
)
def combine(
self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
) -> torch.Tensor:
"""Combine the hidden states and router logits from the appropriate device.
This is a no-op in the base class.
"""
assert self.all2all_manager is not None
return self.all2all_manager.combine(
hidden_states,
is_sequence_parallel,
)
class _CPUSHMDistributed:
def __init__(self, communicator: CpuCommunicator):
self.communicator = communicator
self.group_name = self.make_group_name(communicator)
self.handle = self._init_cpu_shm()
@staticmethod
def make_group_name(communicator: CpuCommunicator) -> str:
instance_identifier = os.environ["VLLM_DIST_IDENT"]
unique_name = communicator.unique_name
instance_identifier = f"{instance_identifier}-{unique_name}"
group_ranks = [str(rank) for rank in communicator.ranks]
shm_group_identifier = f"[{'-'.join(group_ranks)}]"
return f"{instance_identifier}-{shm_group_identifier}-cpushm"
def _init_cpu_shm(self) -> int:
thread_num_tensor = torch.tensor(
[torch.get_num_threads()],
dtype=torch.int64,
)
torch.distributed.all_reduce(
thread_num_tensor,
op=torch.distributed.ReduceOp.MIN,
group=self.communicator.device_group,
)
thread_num = thread_num_tensor.item()
handle = torch.ops._C.init_shm_manager(
self.group_name,
self.communicator.world_size,
self.communicator.rank,
thread_num,
)
torch.distributed.barrier(self.communicator.device_group)
torch.ops._C.join_shm_manager(
handle,
self.group_name,
)
torch.distributed.barrier(self.communicator.device_group)
return handle
def all_reduce(
self, input: torch.Tensor, group: ProcessGroup | None = None
) -> None:
torch.ops._C.shm_allreduce(self.handle, input)
def gather(
self,
input: torch.Tensor,
gather_list: list[torch.Tensor] | None,
dst: int = -1,
group: ProcessGroup | None = None,
) -> None:
# Note: different from the torch gather, here we use local dst rank.
torch.ops._C.shm_gather(
self.handle,
input,
gather_list,
torch.distributed.get_group_rank(group, dst),
)
def all_gather_into_tensor(
self,
output: torch.Tensor,
input: torch.Tensor,
group: ProcessGroup | None = None,
) -> None:
torch.ops._C.shm_all_gather(self.handle, input, output)
def send_tensor_dict(
self,
tensor_dict: dict[str, torch.Tensor | Any],
dst: int,
) -> None:
key_list = list(tensor_dict.keys())
value_list = list(tensor_dict.values())
size_list = []
for v in value_list:
if not isinstance(v, torch.Tensor):
raise RuntimeError("CpuCommunicator only supports sending tensors.")
size_list.append(v.size())
key_size_tensor = torch.frombuffer(
pickle.dumps([key_list, size_list]), dtype=torch.uint8
)
value_list.append(key_size_tensor)
torch.ops._C.shm_send_tensor_list(self.handle, value_list, dst)
return None
def recv_tensor_dict(
self,
src: int,
) -> dict[str, torch.Tensor | Any]:
tensor_list = torch.ops._C.shm_recv_tensor_list(self.handle, src)
value_list: list[torch.Tensor] = tensor_list[:-1]
key_size_tensor = tensor_list[-1]
key_size = pickle.loads(key_size_tensor.numpy().tobytes())
key_list = key_size[0]
size_list = key_size[1]
assert len(key_list) == len(size_list)
assert len(key_list) == len(value_list)
tensor_dict: dict[str, torch.Tensor] = {}
for key, size, t in zip(key_list, size_list, value_list):
tensor_dict[key] = t.view(size)
return tensor_dict
vllm/distributed/ec_transfer/ec_connector/cpu/__init__.py
View File
vllm/distributed/ec_transfer/ec_connector/cpu/common.py
-98
View File
@@ -1,98 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared types for the ECCPUConnector scheduler and worker delegates."""
from dataclasses import dataclass, field
from typing import TYPE_CHECKING
import torch
from vllm.distributed.ec_transfer.ec_connector.base import (
ECConnectorMetadata,
ECConnectorWorkerMetadata,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.ec_shared_region import (
ECSharedRegion,
)
if TYPE_CHECKING:
from vllm.config import VllmConfig
@dataclass
class ECCPUConnectorMetadata(ECConnectorMetadata):
"""Per-step scheduler → worker payload for the ECCPUConnector.
Populated by `ECCPUScheduler.build_connector_meta`; consumed by
`ECCPUWorker` via the mixin's `bind_connector_metadata`.
"""
# Producer role: mm_hashes the scheduler has just allocated CPU
# blocks for this step; the worker's save_caches copies
# encoder_cache[mm_hash] → mmap at these block IDs.
saves: dict[str, list[int]] = field(default_factory=dict)
# Consumer role: mm_hashes whose bytes are available in the local mmap,
# mapped to (transfer_id, block_ids, shape); the worker's start_load_caches
# copies mmap[block_ids] → GPU encoder_cache and reports the transfer_id
# once its copy lands. The id is minted per dispatch, so two loads of the
# same mm_hash in different steps are distinct and a late report from the
# earlier one cannot release the later one's pin.
loads: dict[str, tuple[int, list[int], tuple[int, int]]] = field(
default_factory=dict
)
@dataclass
class ECCPUWorkerMetadata(ECConnectorWorkerMetadata):
"""Per-step worker → scheduler payload for the ECCPUConnector.
Reports the GPU copies that completed this step: saved mm_hashes become
safe to mark ready, and loaded transfers become safe to unpin once every
participating rank has reported them. Built by
`ECCPUWorker.build_connector_worker_meta`; consumed by
`ECCPUScheduler.update_connector_output`.
Loads are reported by transfer id rather than mm_hash because every rank
copies the same blocks: `aggregate` concatenates, so an id appears once
per reporting rank and the scheduler can count participants off the list.
"""
completed_saves: list[str] = field(default_factory=list)
completed_loads: list[int] = field(default_factory=list)
def aggregate(self, other: ECConnectorWorkerMetadata) -> ECConnectorWorkerMetadata:
assert isinstance(other, ECCPUWorkerMetadata)
self.completed_saves.extend(other.completed_saves)
self.completed_loads.extend(other.completed_loads)
return self
def create_ec_shared_region(vllm_config: "VllmConfig") -> ECSharedRegion:
"""Build the EC mmap region from `vllm_config`.
Both `ECCPUScheduler` and `ECCPUWorker` call this to get the same
shared region (same engine_id, same block_size_bytes).
"""
ec_config = vllm_config.ec_transfer_config
assert ec_config is not None, "ec_transfer_config required to build region"
dp_index = vllm_config.parallel_config.data_parallel_index
engine_id = f"{vllm_config.instance_id}_dp{dp_index}"
dtype = vllm_config.model_config.dtype
hidden_dim = vllm_config.model_config.get_inputs_embeds_size()
element_size = torch.empty(0, dtype=dtype).element_size()
block_size_bytes = hidden_dim * element_size
cpu_bytes = ec_config.ec_connector_extra_config.get("ec_cpu_bytes")
if not cpu_bytes:
raise ValueError("ec_cpu_bytes must be specified in ec_connector_extra_config")
cpu_bytes = int(cpu_bytes)
num_blocks = cpu_bytes // block_size_bytes
return ECSharedRegion(
engine_id=engine_id,
num_blocks=num_blocks,
block_size_bytes=block_size_bytes,
)
vllm/distributed/ec_transfer/ec_connector/cpu/connector.py
-153
View File
@@ -1,153 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""ECCPUConnector — CPU encoder-cache offloading.
A role-routed shell: one instance per process. The scheduler delegate owns the
mmap region and offload bookkeeping; the worker delegate owns the mmap view and
GPU copy plumbing. An ec_both instance reuses encoder outputs it has already
offloaded to CPU instead of recomputing them.
"""
from typing import TYPE_CHECKING, Any
import torch
from vllm.distributed.ec_transfer.ec_connector.base import (
ECConnectorBase,
ECConnectorRole,
ECConnectorWorkerMetadata,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUConnectorMetadata,
)
from vllm.logger import init_logger
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ECConnectorOutput
from vllm.v1.request import Request
logger = init_logger(__name__)
class ECCPUConnector(ECConnectorBase):
"""EC connector that offloads encoder cache to a shared CPU mmap region."""
def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole) -> None:
super().__init__(vllm_config=vllm_config, role=role)
if not vllm_config.use_v2_model_runner:
raise ValueError(
"ECCPUConnector requires the V2 model runner. "
"Set VLLM_USE_V2_MODEL_RUNNER=1."
)
self.connector_worker = None
self.connector_scheduler = None
if role == ECConnectorRole.WORKER:
self.connector_worker = self._make_worker(vllm_config)
elif role == ECConnectorRole.SCHEDULER:
self.connector_scheduler = self._make_scheduler(vllm_config)
else:
raise ValueError(f"Unknown ECConnectorRole: {role}")
# Construction seams.
def _make_worker(self, vllm_config: "VllmConfig"):
# Deferred import: the worker module touches torch/CUDA at import time
# via the region, so keep that cost off the scheduler path.
from vllm.distributed.ec_transfer.ec_connector.cpu.worker import (
ECCPUWorker,
)
return ECCPUWorker(vllm_config)
def _make_scheduler(self, vllm_config: "VllmConfig"):
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler import (
ECCPUScheduler,
)
return ECCPUScheduler(vllm_config)
# Worker-side forwarders.
def start_load_caches(
self, encoder_cache: dict[str, torch.Tensor], **kwargs
) -> None:
assert self.connector_worker is not None
metadata = self._get_connector_metadata()
assert isinstance(metadata, ECCPUConnectorMetadata)
self.connector_worker.start_load_caches(
encoder_cache, connector_metadata=metadata
)
def save_caches(
self, encoder_cache: dict[str, torch.Tensor], mm_hash: str, **kwargs
) -> None:
assert self.connector_worker is not None
metadata = self._get_connector_metadata()
assert isinstance(metadata, ECCPUConnectorMetadata)
self.connector_worker.save_caches(
encoder_cache, mm_hash, connector_metadata=metadata
)
# Scheduler-side forwarders.
def has_cache_item(self, identifier: str) -> bool:
assert self.connector_scheduler is not None
return self.connector_scheduler.has_cache_item(identifier)
def ensure_cache_available(
self, request: "Request", num_computed_tokens: int
) -> bool:
assert self.connector_scheduler is not None
return self.connector_scheduler.ensure_cache_available(
request, num_computed_tokens
)
def take_unavailable_requests(self) -> set[str]:
assert self.connector_scheduler is not None
return self.connector_scheduler.take_unavailable_requests()
def update_state_after_alloc(self, request: "Request", index: int) -> None:
assert self.connector_scheduler is not None
self.connector_scheduler.update_state_after_alloc(request, index)
def build_connector_meta(
self, scheduler_output: "SchedulerOutput"
) -> ECCPUConnectorMetadata:
assert self.connector_scheduler is not None
return self.connector_scheduler.build_connector_meta(scheduler_output)
def request_finished(
self, request: "Request"
) -> tuple[bool, "dict[str, Any] | None"]:
assert self.connector_scheduler is not None
return self.connector_scheduler.request_finished(request)
def get_finished(
self, finished_req_ids: set[str]
) -> tuple[set[str] | None, set[str] | None]:
if self.connector_worker is not None:
self.connector_worker.flush_saves()
return None, None
def build_connector_worker_meta(self) -> ECConnectorWorkerMetadata | None:
if self.connector_worker is not None:
return self.connector_worker.build_connector_worker_meta()
return None
def update_connector_output(self, connector_output: "ECConnectorOutput") -> None:
if self.connector_scheduler is not None:
self.connector_scheduler.update_connector_output(connector_output)
def has_pending_push_work(self) -> bool:
if self.connector_scheduler is not None:
return self.connector_scheduler.has_pending_push_work()
return False
# Shared.
def shutdown(self) -> None:
if self.connector_scheduler is not None:
self.connector_scheduler.shutdown()
if self.connector_worker is not None:
self.connector_worker.shutdown()
vllm/distributed/ec_transfer/ec_connector/cpu/control/__init__.py
View File
vllm/distributed/ec_transfer/ec_connector/cpu/control/base.py
-37
View File
@@ -1,37 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Abstract base classes for the EC control plane.
ControlConnection wraps a single peer channel (one DEALER socket).
"""
from abc import ABC, abstractmethod
class ControlConnection(ABC):
"""Bidirectional message channel to a single remote peer.
All I/O is non-blocking: send enqueues without blocking; recv drains
all buffered inbound messages since the last call.
"""
@property
@abstractmethod
def alive(self) -> bool:
"""True if the connection is still usable."""
@abstractmethod
def send(self, msg: bytes) -> None:
"""Enqueue msg for delivery. Must not block."""
@abstractmethod
def recv(self) -> list[bytes]:
"""Drain and return all buffered inbound messages."""
@abstractmethod
def mark_dead(self) -> None:
"""Signal that the peer disconnected. Sets alive to False."""
@abstractmethod
def close(self) -> None:
"""Release all resources. Idempotent."""
vllm/distributed/ec_transfer/ec_connector/cpu/control/zmq.py
-233
View File
@@ -1,233 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""ZMQ control-plane implementations for ECCPUConnector.
ZmqClientConnection — one DEALER socket to a single producer peer.
ZmqClientTransport — pool of ZmqClientConnection objects (consumer side).
ZmqServerTransport — ROUTER socket (producer side).
"""
import contextlib
from dataclasses import dataclass, field
import zmq
from zmq.utils.monitor import recv_monitor_message
from vllm.distributed.ec_transfer.ec_connector.cpu.control.base import (
ControlConnection,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr
from vllm.logger import init_logger
from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
logger = init_logger(__name__)
_HEARTBEAT_IVL_MS = 2000
_HEARTBEAT_TIMEOUT_MS = 4000
_HEARTBEAT_TTL_MS = 8000
# ── Client side ───────────────────────────────────────────────────────────────
@dataclass
class ZmqClientConnection(ControlConnection):
"""One DEALER socket to one producer peer.
Owns the socket and its optional monitor socket. NIXL state is NOT
stored here — it belongs in ECSession above this layer.
"""
dealer: zmq.Socket
monitor: zmq.Socket | None
_alive: bool = field(default=True, init=False)
@property
def alive(self) -> bool:
return self._alive
def send(self, msg: bytes) -> None:
self.dealer.send_multipart([b"", msg])
def recv(self) -> list[bytes]:
msgs: list[bytes] = []
while True:
try:
frames = self.dealer.recv_multipart(flags=zmq.NOBLOCK)
except zmq.Again:
break
except Exception:
logger.exception("ec: DEALER recv failed")
break
if len(frames) != 2 or frames[0] != b"":
logger.warning(
"ec: dropped malformed frame envelope (%d frames)", len(frames)
)
continue
msgs.append(frames[1])
return msgs
def mark_dead(self) -> None:
self._alive = False
def close(self) -> None:
if self.monitor is not None:
with contextlib.suppress(Exception):
self.dealer.disable_monitor()
with contextlib.suppress(Exception):
self.monitor.close(linger=0)
with contextlib.suppress(Exception):
self.dealer.close(linger=0)
self._alive = False
class ZmqClientTransport:
"""Pool of ZmqClientConnection objects — one per producer peer.
Creates DEALER sockets lazily on first contact. Does not touch NIXL;
NIXL registration is handled by ECCPUConsumer once an XferAck arrives.
"""
def __init__(self) -> None:
self._ctx = zmq.Context()
self._connections: dict[PeerAddr, ZmqClientConnection] = {}
def connect(self, addr: PeerAddr) -> ZmqClientConnection:
"""Return the existing connection for addr, or create a new one."""
existing = self._connections.get(addr)
if existing is not None:
return existing
host, port = addr
dealer = make_zmq_socket(
ctx=self._ctx,
path=make_zmq_path(scheme="tcp", host=host, port=port),
socket_type=zmq.DEALER,
bind=False,
)
dealer.setsockopt(zmq.HEARTBEAT_IVL, _HEARTBEAT_IVL_MS)
dealer.setsockopt(zmq.HEARTBEAT_TIMEOUT, _HEARTBEAT_TIMEOUT_MS)
dealer.setsockopt(zmq.HEARTBEAT_TTL, _HEARTBEAT_TTL_MS)
monitor_addr = f"inproc://ec-peer-mon-{host}-{port}"
dealer.monitor(monitor_addr, zmq.EVENT_DISCONNECTED)
monitor = self._ctx.socket(zmq.PAIR)
monitor.connect(monitor_addr)
conn = ZmqClientConnection(dealer=dealer, monitor=monitor)
self._connections[addr] = conn
return conn
def poll(self) -> dict[PeerAddr, list[bytes]]:
"""Non-blocking drain of messages from every peer's DEALER."""
msgs: dict[PeerAddr, list[bytes]] = {}
for addr, conn in list(self._connections.items()):
conn_msgs = conn.recv()
msgs[addr] = conn_msgs
return msgs
def poll_dead(self) -> list[PeerAddr]:
"""Non-blocking sweep of peer monitors; return addresses that died."""
dead: list[PeerAddr] = []
for addr, conn in list(self._connections.items()):
if conn.monitor is None:
continue
try:
while True:
evt = recv_monitor_message(conn.monitor, flags=zmq.NOBLOCK)
if evt["event"] == zmq.EVENT_DISCONNECTED:
conn.mark_dead()
dead.append(addr)
break
except zmq.Again:
pass
except Exception:
logger.warning("ec: monitor poll failed for %s", addr, exc_info=True)
return dead
def remove(self, addr: PeerAddr) -> ZmqClientConnection | None:
"""Remove and close the connection for addr, returning it (or None)."""
conn = self._connections.pop(addr, None)
if conn is not None:
conn.close()
return conn
def close(self) -> None:
for conn in list(self._connections.values()):
conn.close()
self._connections.clear()
try:
self._ctx.destroy(linger=0)
except Exception:
logger.debug("ec: zmq context destroy failed", exc_info=True)
# ── Server side ───────────────────────────────────────────────────────────────
class ZmqServerTransport:
"""ROUTER socket for the producer side. Byte-level only, driven by
ProducerSession.
poll(timeout_ms) waits up to timeout_ms for messages, then drains all that
are available; timeout_ms=0 drains whatever has already arrived. send()
routes a reply back to a specific peer by ZMQ identity.
"""
def __init__(self, host: str, port: int) -> None:
self._ctx = zmq.Context()
self._router = make_zmq_socket(
ctx=self._ctx,
path=make_zmq_path(scheme="tcp", host=host, port=port),
socket_type=zmq.ROUTER,
bind=True,
)
self._poller = zmq.Poller()
self._poller.register(self._router, zmq.POLLIN)
def poll(self, timeout_ms: int = 0) -> list[tuple[bytes, bytes]]:
"""Wait up to timeout_ms for incoming messages, then drain all pending.
Returns a list of (identity, payload) pairs.
"""
if timeout_ms > 0:
try:
events = dict(self._poller.poll(timeout=timeout_ms))
except zmq.ContextTerminated:
return []
except Exception:
logger.exception("ec: router poller failed")
return []
if self._router not in events:
return []
msgs: list[tuple[bytes, bytes]] = []
while True:
try:
identity, _, payload = self._router.recv_multipart(flags=zmq.NOBLOCK)
msgs.append((identity, payload))
except zmq.Again:
break
except zmq.ContextTerminated:
break
except Exception:
logger.exception("ec: router recv failed")
break
return msgs
def send(self, identity: bytes, payload: bytes) -> None:
"""Send a reply payload to the peer identified by ZMQ identity."""
try:
self._router.send_multipart([identity, b"", payload])
except Exception:
logger.exception("ec: router send failed")
def close(self) -> None:
try:
self._router.close(linger=0)
except Exception:
logger.debug("ec: router close failed", exc_info=True)
try:
self._ctx.destroy(linger=0)
except Exception:
logger.debug("ec: zmq context destroy failed", exc_info=True)
vllm/distributed/ec_transfer/ec_connector/cpu/data/__init__.py
View File
vllm/distributed/ec_transfer/ec_connector/cpu/data/base.py
-67
View File
@@ -1,67 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Abstract base class for the EC data-plane transport (NIXL/UCX)."""
from abc import ABC, abstractmethod
from typing import Any
class DataTransport(ABC):
"""Interface wrapping the NIXL data-plane for EC block transfers.
Peer lifecycle:
add_remote_peer() — register a peer and store its dlist handle internally
remove_remote_peer() — deregister and release the stored handle
post_read() — issue a READ using the stored handle for agent_name
The dlist handle (prepped xfer descriptor list) is an implementation detail
of the transport and is never exposed to callers.
"""
@abstractmethod
def get_agent_metadata(self) -> bytes:
"""Return this agent's serialized NIXL metadata blob."""
@abstractmethod
def get_mem_descriptor(self) -> bytes:
"""Return the msgpack-encoded block descriptor list for this agent."""
@abstractmethod
def add_remote_peer(self, metadata: bytes, mem_descriptor: bytes) -> str:
"""Register a remote peer, prep its READ-source dlist, and return agent_name.
The dlist handle is stored internally and used by post_read().
"""
@abstractmethod
def remove_remote_peer(self, agent_name: str) -> None:
"""Deregister a previously registered remote peer and release its handle."""
@abstractmethod
def post_read(
self,
local_indices: list[int],
agent_name: str,
remote_indices: list[int],
notif_msg: bytes,
) -> Any:
"""Issue a consumer-initiated READ using the stored dlist for agent_name.
Returns a transfer handle for use with check_xfer_state / release_xfer_handle.
"""
@abstractmethod
def get_new_notifs(self) -> dict[str, list[bytes]]:
"""Drain NIXL completion notifications addressed to this agent."""
@abstractmethod
def check_xfer_state(self, handle: Any) -> str:
"""Return the NIXL state string: 'DONE', 'PROC', or an error string."""
@abstractmethod
def release_xfer_handle(self, handle: Any) -> None:
"""Release a transfer handle returned by post_read."""
@abstractmethod
def deregister(self) -> None:
"""Deregister all memory and tear down the NIXL agent."""
vllm/distributed/ec_transfer/ec_connector/cpu/data/nixl.py
-162
View File
@@ -1,162 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""NIXL-backed data-plane transport for ECCPUConnector."""
from typing import Any
from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import (
DataTransport,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.utils import (
build_block_descs,
deserialize_mem_descriptor,
serialize_mem_descriptor,
)
from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
from vllm.logger import init_logger
logger = init_logger(__name__)
_NIXL_DRAM = "DRAM"
class NixlDataTransport(DataTransport):
"""Adapts NixlWrapper to the DataTransport interface.
Dlist handles (prepped xfer descriptor lists for remote peers) are stored
internally keyed by agent_name. Callers never touch raw NIXL handles —
post_read() resolves the handle from agent_name internally.
NixlWrapper is not thread-safe, and neither is this class. Every call site
runs on the scheduler thread, which serializes the transport by
construction; transfers are async inside NIXL and are drained from that
same thread via get_new_notifs() / check_xfer_state().
"""
def __init__(
self,
agent_name: str,
base_ptr: int,
num_blocks: int,
block_size_bytes: int,
total_size_bytes: int,
num_threads: int = 1,
) -> None:
if NixlWrapper is None or nixl_agent_config is None:
raise RuntimeError(
"ECCPUConnector requires NIXL; "
"install the `nixl` package or set a different ec_connector."
)
# agent_name → remote dlist handle (from add_remote_peer, used by post_read)
self._peer_handles: dict[str, int] = {}
# None until the local region is registered, so deregister() is safe on
# a transport whose registration failed part-way.
self._reg_descs: Any = None
self._nixl = NixlWrapper(
agent_name,
nixl_agent_config(num_threads=num_threads, capture_telemetry=True),
)
self._register_local_memory(
base_ptr, num_blocks, block_size_bytes, total_size_bytes
)
def _register_local_memory(
self,
base_ptr: int,
num_blocks: int,
block_size_bytes: int,
total_size_bytes: int,
) -> None:
"""Register the local region with NIXL and prep its transfer dlist.
Unwinds the registration if a later step fails: __init__ propagates the
error, so the caller never receives an object it could deregister with.
"""
block_descs = build_block_descs(base_ptr, num_blocks, block_size_bytes)
self._reg_descs = self._nixl.get_reg_descs(
[(base_ptr, total_size_bytes, 0, "")], _NIXL_DRAM
)
self._nixl.register_memory(self._reg_descs, backends=["UCX"])
try:
xfer_descs = self._nixl.get_xfer_descs(block_descs, _NIXL_DRAM)
self._local_xfer_handle = self._nixl.prep_xfer_dlist(
"NIXL_INIT_AGENT", xfer_descs
)
self._agent_metadata = self._nixl.get_agent_metadata()
self._mem_descriptor_bytes = serialize_mem_descriptor(block_descs)
except Exception:
self.deregister()
raise
def get_agent_metadata(self) -> bytes:
return self._agent_metadata
def get_mem_descriptor(self) -> bytes:
return self._mem_descriptor_bytes
def add_remote_peer(self, metadata: bytes, mem_descriptor: bytes) -> str:
"""Register remote peer, prep its dlist, store handle internally.
Returns agent_name for use in post_read() and remove_remote_peer().
"""
agent_name = self._nixl.add_remote_agent(metadata)
remote_blocks = deserialize_mem_descriptor(mem_descriptor)
remote_xfer_descs = self._nixl.get_xfer_descs(remote_blocks, _NIXL_DRAM)
self._peer_handles[agent_name] = self._nixl.prep_xfer_dlist(
agent_name, remote_xfer_descs
)
return agent_name
def remove_remote_peer(self, agent_name: str) -> None:
self._peer_handles.pop(agent_name, None)
try:
self._nixl.remove_remote_agent(agent_name)
except Exception:
logger.warning(
"EC: remove_remote_peer failed for %s", agent_name, exc_info=True
)
def post_read(
self,
local_indices: list[int],
agent_name: str,
remote_indices: list[int],
notif_msg: bytes,
) -> Any:
if len(local_indices) != len(remote_indices):
raise ValueError(
f"EC: local/remote block count mismatch "
f"({len(local_indices)} vs {len(remote_indices)})"
)
remote_handle = self._peer_handles[agent_name]
handle = self._nixl.make_prepped_xfer(
"READ",
self._local_xfer_handle,
local_indices,
remote_handle,
remote_indices,
notif_msg=notif_msg,
)
self._nixl.transfer(handle)
return handle
def get_new_notifs(self) -> dict[str, list[bytes]]:
return self._nixl.get_new_notifs()
def check_xfer_state(self, handle: Any) -> str:
return self._nixl.check_xfer_state(handle)
def release_xfer_handle(self, handle: Any) -> None:
try:
self._nixl.release_xfer_handle(handle)
except Exception:
logger.warning("EC: release_xfer_handle failed", exc_info=True)
def deregister(self) -> None:
if self._reg_descs is None:
return
try:
self._nixl.deregister_memory(self._reg_descs)
except Exception:
logger.warning("EC: deregister failed", exc_info=True)
self._reg_descs = None
vllm/distributed/ec_transfer/ec_connector/cpu/ec_shared_region.py
-193
View File
@@ -1,193 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Lightweight mmap-backed shared memory region for encoder cache (EC) data.
Modeled after SharedOffloadRegion (vllm/v1/kv_offload/cpu/) but simplified
for EC: flat shared layout, no multi-tensor cursor, no block_size_factor.
"""
import errno
import mmap
import os
import time
from collections.abc import Callable
import numpy as np
import torch
from vllm.logger import init_logger
logger = init_logger(__name__)
# MADV_POPULATE_WRITE was added in Linux 5.14 (value 23).
_MADV_POPULATE_WRITE = getattr(mmap, "MADV_POPULATE_WRITE", 23)
def _madvise_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None:
mmap_obj.madvise(_MADV_POPULATE_WRITE, offset, length)
def _fallback_populate_write(mmap_obj: mmap.mmap, offset: int, length: int) -> None:
# Touch one byte per page via a read-modify-write so existing bytes are
# preserved — a peer worker may have already written EC data into this
# shared mmap by the time we run on a kernel without MADV_POPULATE_WRITE.
arr = np.frombuffer(mmap_obj, dtype=np.uint8)
arr[offset : offset + length : mmap.PAGESIZE] |= 0
def _get_populate_write_fn(
mmap_obj: mmap.mmap,
) -> Callable[[mmap.mmap, int, int], None]:
"""Select the pre-faulting method once for this mmap."""
try:
_madvise_populate_write(mmap_obj, 0, mmap.PAGESIZE)
except OSError as e:
if e.errno != errno.EINVAL:
raise
logger.warning(
"MADV_POPULATE_WRITE is not supported; falling back to per-page "
"writes for EC mmap pre-population. Startup may be slower."
)
return _fallback_populate_write
return _madvise_populate_write
def _wait_for_file_size(fd: int, expected_size: int, timeout: float = 30.0):
"""Spin-wait until the file reaches expected_size (creator truncated it)."""
deadline = time.monotonic() + timeout
while True:
if os.fstat(fd).st_size >= expected_size:
return
if time.monotonic() > deadline:
raise TimeoutError(
f"Timed out waiting for EC mmap file to reach {expected_size} bytes"
)
time.sleep(0.005)
class ECSharedRegion:
"""Flat mmap-backed memory region shared across TP workers for
encoder cache blocks.
Layout: (num_blocks, block_size_bytes) — contiguous, no per-worker
interleaving. All workers map the same file and see identical data.
File path: /dev/shm/vllm_ec_{engine_id}.mmap
This class owns only the shared memory substrate (mmap lifecycle, the
`blocks` view, CUDA host registration). Block allocation and eviction
are tracked by `EmbeddingCache` in the scheduler process.
"""
def __init__(
self,
engine_id: str,
num_blocks: int,
block_size_bytes: int,
) -> None:
self.num_blocks = num_blocks
self.block_size_bytes = block_size_bytes
total_size_bytes = num_blocks * block_size_bytes
# Path in /dev/shm (tmpfs); unique per engine instance.
self._mmap_path = f"/dev/shm/vllm_ec_{engine_id}.mmap"
# True for the process that created the file (responsible for unlink).
self._is_creator = False
# True after successful cudaHostRegister (cleanup must unregister).
self._is_pinned = False
# File descriptor for the shared memory backing file.
try:
self._fd: int | None = os.open(
self._mmap_path, os.O_CREAT | os.O_EXCL | os.O_RDWR, 0o600
)
os.ftruncate(self._fd, total_size_bytes)
self._is_creator = True
logger.info(
"Created EC mmap file %s (%.2f MiB)",
self._mmap_path,
total_size_bytes / (1 << 20),
)
except FileExistsError:
self._fd = os.open(self._mmap_path, os.O_RDWR)
try:
_wait_for_file_size(self._fd, total_size_bytes)
except Exception:
os.close(self._fd)
self._fd = None
raise
logger.info("Opened existing EC mmap file %s", self._mmap_path)
# MAP_SHARED mmap over _fd; all processes see the same pages.
self._mmap_obj: mmap.mmap | None = mmap.mmap(
self._fd,
total_size_bytes,
flags=mmap.MAP_SHARED,
prot=mmap.PROT_READ | mmap.PROT_WRITE,
)
if self._is_creator:
populate_write_fn = _get_populate_write_fn(self._mmap_obj)
populate_write_fn(self._mmap_obj, 0, total_size_bytes)
# (num_blocks, block_size_bytes) int8 tensor over the mmap buffer.
self.blocks: torch.Tensor = torch.frombuffer(
memoryview(self._mmap_obj), dtype=torch.int8
).view(num_blocks, block_size_bytes)
# Cached for cudaHostRegister/Unregister and pointer math.
self._blocks_ptr: int = self.blocks.data_ptr()
self._blocks_nbytes: int = self.blocks.nbytes
def pin_memory(self) -> None:
"""Register the entire mmap as CUDA pinned memory for fast DMA.
Each TP worker process owns its own virtual address for the same
shared physical pages, so every process must register independently.
No-op when CUDA is not available or already pinned.
"""
if self._is_pinned or not torch.cuda.is_available():
return
result = torch.cuda.cudart().cudaHostRegister(
self._blocks_ptr, self._blocks_nbytes, 0
)
if result.value != 0:
logger.warning(
"cudaHostRegister failed (code=%d) — "
"transfers will still work but may be slower (unpinned DMA)",
result.value,
)
else:
logger.debug("cudaHostRegister %.2f MB", self._blocks_nbytes / 1e6)
self._is_pinned = True
def cleanup(self) -> None:
"""Tear down the region. Lifecycle method; no concurrent access."""
logger.info("Starting ECSharedRegion cleanup...")
if self._is_creator:
try:
os.unlink(self._mmap_path)
logger.info("Removed EC mmap file %s", self._mmap_path)
except Exception:
logger.warning(
"Failed to unlink path %s", self._mmap_path, exc_info=True
)
self._is_creator = False
if self._is_pinned:
result = torch.cuda.cudart().cudaHostUnregister(self._blocks_ptr)
if result.value != 0:
logger.warning("cudaHostUnregister failed (code=%d)", result)
self._is_pinned = False
if hasattr(self, "blocks"):
del self.blocks
if self._mmap_obj:
try:
self._mmap_obj.close()
except Exception:
logger.warning("Failed to close mmap_obj", exc_info=True)
self._mmap_obj = None
if self._fd is not None:
try:
os.close(self._fd)
except Exception:
logger.warning("Failed to close fd %s", self._fd, exc_info=True)
self._fd = None
vllm/distributed/ec_transfer/ec_connector/cpu/protocol.py
-91
View File
@@ -1,91 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Scheduler-side wire types for the ECCPUConnector.
Two peer-to-peer message types exchanged between consumer and producer over
ZMQ (XferReq, XferAck) plus a peer-compatibility fingerprint
(compute_ec_compatibility_hash) sent on every XferReq so a producer refuses
to serve a mismatched peer.
"""
import enum
import hashlib
import json
import msgspec
# Bump when the wire format or tensor/block layout changes in a
# backward-incompatible way. Peers reject mismatched versions.
EC_CONNECTOR_VERSION: int = 2
class XferStatus(enum.IntEnum):
"""Outcome the producer reports for an XferReq."""
OK = 0
NACK_MISSING = 1 # producer no longer holds the encoding
NACK_INCOMPAT = 2 # peer-compatibility hash mismatch
NACK_VERSION = 3 # connector wire-version mismatch
NACK_INTERNAL = 4 # producer hit an unexpected error
NACK_NOT_READY = 5 # announced, but its GPU->mmap save has not landed yet
# The producer expects a retryable condition to clear on its own, so the
# consumer re-requests the read rather than giving up on the encoding.
RETRYABLE_NACKS = frozenset((XferStatus.NACK_NOT_READY,))
# Ordinary outcomes rather than signs of a broken peer.
EXPECTED_NACKS = RETRYABLE_NACKS | {XferStatus.NACK_MISSING}
class XferReq( # type: ignore[call-arg]
msgspec.Struct,
tag="req",
omit_defaults=True,
):
"""Consumer → Producer: I want to READ mm_hash; pin it and tell me where."""
mm_hash: str
compatibility_hash: str
# Consumer's session identity. Together with mm_hash it uniquely identifies
# this transfer on both sides: the producer keys its active xfers by
# (session_id, mm_hash), and the NIXL notif_msg carries the same pair.
session_id: str = ""
connector_version: int = EC_CONNECTOR_VERSION
class XferAck( # type: ignore[call-arg]
msgspec.Struct,
tag="ack",
omit_defaults=True,
):
"""Producer → Consumer: response to an XferReq.
On status == OK: producer has pinned source blocks and returns their
indices plus fresh NIXL metadata for the consumer to register and READ.
On any NACK_*: optional fields are empty; consumer falls back to local encode.
"""
mm_hash: str
status: XferStatus
# Consumer's session_id echoed back from XferReq. Together with mm_hash it
# identifies the exact transfer; the consumer uses both to build notif_msg.
session_id: str = ""
src_block_indices: list[int] = []
agent_metadata: bytes = b""
mem_descriptor: bytes = b""
def compute_ec_compatibility_hash(
vllm_version: str,
model: str,
dtype: str,
block_size_bytes: int,
) -> str:
"""Peer-compatibility fingerprint over the four factors that determine
byte-layout compatibility. Producer NACKs any XferReq whose hash differs.
"""
factors = json.dumps(
[vllm_version, model, dtype, block_size_bytes],
separators=(",", ":"),
)
return hashlib.sha256(factors.encode("utf-8")).hexdigest()
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/__init__.py
-700
View File
@@ -1,700 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""ECCPUScheduler — CPU offload scheduler delegate.
Owns the mmap region and the embedding cache, and handles the producer
(GPU->CPU offload) and consumer (CPU->GPU reload) scheduler-side logic
for the ECCPUConnector.
"""
import time
from collections.abc import Collection
from typing import TYPE_CHECKING, Any
import torch
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUConnectorMetadata,
ECCPUWorkerMetadata,
create_ec_shared_region,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
EmbeddingCache,
)
from vllm.distributed.ec_transfer.ec_connector.encoder_output_width import (
get_encoder_output_width,
)
from vllm.distributed.ec_transfer.ec_connector.utils import (
PlaceholderMetadataResolver,
collect_ec_item_metadata,
)
from vllm.logger import init_logger
if TYPE_CHECKING:
from vllm.config import VllmConfig
from vllm.multimodal.inputs import MultiModalFeatureSpec
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.outputs import ECConnectorOutput
from vllm.v1.request import Request
logger = init_logger(__name__)
# Total remote wait budget per request item, including all retries.
_ADMIT_DEFER_TIMEOUT_S = 60.0
class _RemoteUnavailable(Exception):
"""An item cannot be fetched within its request's remote wait budget."""
class ECCPUScheduler:
"""Scheduler delegate for the ECCPUConnector."""
def __init__(self, vllm_config: "VllmConfig") -> None:
self._vllm_config = vllm_config
ec_config = vllm_config.ec_transfer_config
assert ec_config is not None
self._is_producer: bool = ec_config.is_ec_producer
self._is_consumer: bool = ec_config.is_ec_consumer
self._region = create_ec_shared_region(vllm_config)
# Block allocator + LRU eviction policy for the shared region.
self._cache = EmbeddingCache(self._region.num_blocks)
self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config)
# mm_hash → block IDs allocated this step for GPU→mmap saves.
self._pending_saves: dict[str, list[int]] = {}
# mm_hash → (transfer_id, block IDs, shape) to load from mmap→GPU.
self._pending_loads: dict[str, tuple[int, list[int], tuple[int, int]]] = {}
# Dispatched loads awaiting completion reports, keyed by transfer id:
# transfer_id → (mm_hash, reports still outstanding). The pin taken at
# dispatch is released only once the count reaches zero.
self._load_acks: dict[int, tuple[str, int]] = {}
self._next_transfer_id = 0
pc = vllm_config.parallel_config
# Reports to expect per load. Only pipeline stage 0 runs the EC
# connector, so the participants are the tp × pcp ranks of that stage.
# Other executors deliver a single rank's output, and expecting reports
# that cannot arrive would hold the pin forever.
self._expected_load_reports = (
pc.tensor_parallel_size * pc.prefill_context_parallel_size
if pc.distributed_executor_backend == "mp"
else 1
)
self._ec_config = ec_config
# NIXL p2p is an option of this connector, not engine-wide behavior, so
# it lives in extra config. Extra config is not type coerced, so a value
# that arrived as a JSON/CLI string is parsed, not truth-tested.
raw_nixl = ec_config.get_from_extra_config("ec_enable_nixl", False)
self._nixl_enabled = (
raw_nixl
if isinstance(raw_nixl, bool)
else str(raw_nixl).strip().lower() in ("true", "1", "yes")
)
# NIXL fields default to None/empty so the gate-off path is untouched.
self._data: Any = None
self._compat_hash: str | None = None
self._first_in_batch = True
self._transport: Any = None
self._producer_session: Any = None
self._sessions: dict = {}
self._in_flight: set[str] = set()
self._tombstones: set[str] = set()
self._step_completed: set[str] = set()
# Per-request clocks survive retries and are cleared on completion.
self._deferred_since: dict[tuple[str, str], float] = {}
# Requests needing a remote encoding that will not arrive. Drained by
# take_unavailable_requests(); the scheduler aborts them.
self._unrecoverable: set[str] = set()
self._peer_host: str | None = None
self._peer_port: int | None = None
dtype = vllm_config.model_config.dtype
assert isinstance(dtype, torch.dtype)
self._dtype = dtype
self._ack_timeout_s: float = 0.0
if self._nixl_enabled:
try:
self._setup_nixl(vllm_config)
except Exception:
self.shutdown()
raise
def _setup_nixl(self, vllm_config: "VllmConfig") -> None:
# Lazy imports keep nixl/zmq off the gate-off path.
from vllm import envs
from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import (
ZmqClientTransport,
ZmqServerTransport,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.data.nixl import (
NixlDataTransport,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import (
compute_ec_compatibility_hash,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.session import (
PRODUCER_PIN_LEASE_S,
ProducerSession,
)
from vllm.distributed.nixl_utils import NixlWrapper, nixl_agent_config
from vllm.version import __version__ as VLLM_VERSION
if NixlWrapper is None or nixl_agent_config is None:
raise RuntimeError(
"ec_enable_nixl requires NIXL; install the `nixl` package or "
"remove ec_enable_nixl from ec_connector_extra_config."
)
engine_id = self._ec_config.engine_id
assert engine_id is not None
# How long a consumer waits for an XferAck. The producer answers from
# its scheduler step, so its reply latency scales with the encoder's
# batch size: a deployment whose steps run longer must raise this.
self._ack_timeout_s = float(
self._ec_config.get_from_extra_config(
"consumer_ack_timeout_s", PRODUCER_PIN_LEASE_S
)
)
self._compat_hash = compute_ec_compatibility_hash(
vllm_version=VLLM_VERSION,
model=str(vllm_config.model_config.model),
dtype=str(self._dtype),
block_size_bytes=self._region.block_size_bytes,
)
if self._is_producer:
self._peer_host = envs.VLLM_EC_SIDE_CHANNEL_HOST
self._peer_port = envs.VLLM_EC_SIDE_CHANNEL_PORT
self._data = NixlDataTransport(
agent_name=engine_id,
base_ptr=self._region.blocks.data_ptr(),
num_blocks=self._region.num_blocks,
block_size_bytes=self._region.block_size_bytes,
total_size_bytes=self._region.num_blocks * self._region.block_size_bytes,
)
if self._is_producer:
assert self._peer_host is not None
assert self._peer_port is not None
self._producer_session = ProducerSession(
transport=ZmqServerTransport(
host=self._peer_host, port=self._peer_port
),
data=self._data,
cache=self._cache,
compat_hash=self._compat_hash,
)
if self._is_consumer:
self._transport = ZmqClientTransport()
def has_cache_item(self, identifier: str) -> bool:
if not self._is_consumer:
return False
entry = self._cache.get(identifier)
return entry is not None and entry.ready
def ensure_cache_available(
self,
request: "Request",
num_computed_tokens: int,
local_cache_hashes: Collection[str] | None = None,
) -> bool:
if not self._nixl_enabled:
return True # CPU offload never blocks.
first = self._first_in_batch
self._first_in_batch = False
if not self._is_consumer:
return True
if first:
self._poll_step()
return self._nixl_consumer_admit(request, num_computed_tokens)
def _nixl_consumer_admit(
self, request: "Request", num_computed_tokens: int
) -> bool:
"""Admit a request once its required remote encodings are ready.
Each item has a per-request wait budget that survives transfer retries.
Unavailable items fall back to local input; without it, the request is
reported by take_unavailable_requests() for the scheduler to abort.
Returns:
True if no item needs to wait for a remote encoding. False if any
item is pending or the request must be aborted.
"""
if not request.ec_transfer_params:
return True
now = time.monotonic()
pending = False
for feature in request.mm_features:
try:
if not self._admit_item(request, feature, num_computed_tokens, now):
pending = True
except _RemoteUnavailable as error:
if not self._fail_or_fallback(request, feature, str(error)):
return False
return not pending
def _admit_item(
self,
request: "Request",
feature: "MultiModalFeatureSpec",
num_computed_tokens: int,
now: float,
) -> bool:
"""Check one encoding, starting a remote read if needed.
Returns:
True if already computed, cached, or not remotely sourced. False
while waiting for a read or cache space, including across retries.
Raises:
_RemoteUnavailable: The remote source is invalid, the read failed,
or the item's total wait budget expired.
"""
pos = feature.mm_position
mm_hash = feature.identifier
if pos.offset + pos.length <= num_computed_tokens:
self._deferred_since.pop((request.request_id, mm_hash), None)
return True
announced = (request.ec_transfer_params or {}).get(mm_hash)
# Without a producer address there is nothing to fetch: the request
# carries what the model needs and the encoder runs locally.
# `ec_transfer_params` reaches us from the request, so its shape is
# checked rather than assumed.
remote: dict[str, Any] | None = (
announced
if isinstance(announced, dict) and "peer_host" in announced
else None
)
entry = self._cache.get(mm_hash)
if entry is not None and entry.ready:
# Local hit: upstream's update_state_after_alloc pins and
# loads it through the same path as a natively cached entry.
self._deferred_since.pop((request.request_id, mm_hash), None)
return True
if remote is None:
return True
since = self._deferred_since.setdefault((request.request_id, mm_hash), now)
if now - since > _ADMIT_DEFER_TIMEOUT_S:
raise _RemoteUnavailable(f"the remote wait exceeded {now - since:.0f}s")
if mm_hash in self._in_flight or mm_hash in self._step_completed:
return False
if mm_hash in self._tombstones:
self._tombstones.discard(mm_hash)
raise _RemoteUnavailable("the remote read failed")
if entry is not None:
# Present but not ready and not being fetched: its blocks are
# held by a quarantined/settling DMA and cannot be reused.
return False
try:
size = int(remote["size_bytes"])
shape = tuple(int(dim) for dim in remote["shape"])
except (KeyError, TypeError, ValueError) as error:
raise _RemoteUnavailable("the announced size was unusable") from error
# The width is the producer's to report; the row count is not.
num_embeds = pos.get_num_embeds()
if (
len(shape) != 2
or shape[0] != num_embeds
or size != self._num_blocks(shape) * self._region.block_size_bytes
):
logger.warning(
"EC consumer: size mismatch mm_hash=%s announced=%d shape=%s "
"expected %d rows",
mm_hash,
size,
shape,
num_embeds,
)
raise _RemoteUnavailable("the announced size was wrong")
try:
started = self._start_xfer(mm_hash, remote, size, (shape[0], shape[1]))
except Exception as error:
logger.exception(
"EC consumer: failed to start NIXL xfer mm_hash=%s", mm_hash
)
raise _RemoteUnavailable("the read could not be started") from error
if started:
self._in_flight.add(mm_hash)
return False
def _fail_or_fallback(
self, request: "Request", feature: "MultiModalFeatureSpec", why: str
) -> bool:
"""Resolve an unavailable remote encoding and clear its wait budget.
Returns:
True if local media or embeddings are available, removing this
item's remote announcement so subsequent steps do not retry it.
False if only placeholder metadata (or no input) is available,
recording the request for take_unavailable_requests() to drain.
"""
mm_hash = feature.identifier
self._deferred_since.pop((request.request_id, mm_hash), None)
data = feature.data
metadata = (
self._metadata_resolver.fields_for(feature.modality) if data else set()
)
if data and any(
value.data is not None for key, value in data.items() if key not in metadata
):
# Passthrough embeddings bypass the processor cache. SHM address
# items therefore refer to media that the worker can encode.
request.ec_transfer_params = dict(request.ec_transfer_params or {})
request.ec_transfer_params.pop(mm_hash, None)
logger.warning(
"EC consumer: request %s mm_hash=%s: %s; using local input",
request.request_id,
mm_hash,
why,
)
return True
self._unrecoverable.add(request.request_id)
logger.error(
"EC consumer: request %s needs remote encoding mm_hash=%s but %s. "
"No local model input is available; failing the request.",
request.request_id,
mm_hash,
why,
)
return False
def take_unavailable_requests(self) -> set[str]:
"""Return and clear IDs of requests that cannot obtain their encodings.
The scheduler must abort these requests rather than leave them waiting.
Requests that can fall back to local input are not included, and each
recorded failure is returned only once.
"""
if not self._unrecoverable:
return set()
failed = self._unrecoverable
self._unrecoverable = set()
return failed
def _start_xfer(
self,
mm_hash: str,
info: "dict[str, Any]",
size_bytes: int,
shape: tuple[int, int],
) -> bool:
"""Allocate a not-ready cache entry and start a NIXL READ into it.
Returns True if the transfer was started. Returns False when the
cache cannot accommodate the encoding; admission will retry it.
"""
from math import ceil
from vllm.distributed.ec_transfer.ec_connector.cpu.session import (
ConsumerSession,
)
n_blocks = max(1, ceil(size_bytes / self._region.block_size_bytes))
entry = self._cache.alloc(mm_hash, n_blocks, shape)
if entry is None:
logger.debug(
"EC consumer: cache full for mm_hash=%s (%d blocks); deferring",
mm_hash,
n_blocks,
)
return False
indices = list(entry.block_ids)
addr = (info["peer_host"], int(info["peer_port"]))
if addr not in self._sessions:
zmq_conn = self._transport.connect(addr)
assert self._compat_hash is not None
self._sessions[addr] = ConsumerSession(
addr=addr,
zmq_conn=zmq_conn,
transport=self._transport,
data=self._data,
compat_hash=self._compat_hash,
)
deadline = time.monotonic() + self._ack_timeout_s
try:
self._sessions[addr].start_xfer(mm_hash, indices, deadline)
except Exception:
self._cache.discard(mm_hash)
raise
logger.debug(
"EC consumer: starting NIXL xfer mm_hash=%s from %s:%d blocks=%d",
mm_hash,
addr[0],
addr[1],
n_blocks,
)
return True
def _poll_step(self) -> None:
now = time.monotonic()
all_messages = self._transport.poll()
for addr, session in list(self._sessions.items()):
session.poll(all_messages.get(addr, []), now)
for addr in self._transport.poll_dead():
self._on_peer_down(addr)
for session in self._sessions.values():
self._process_session_results(session)
def _process_session_results(self, session: Any) -> None:
r = session.take_results()
for mm_hash in r.completed:
self._in_flight.discard(mm_hash)
self._cache.mark_ready(mm_hash)
self._step_completed.add(mm_hash)
logger.debug("EC consumer: NIXL xfer complete mm_hash=%s", mm_hash)
for mm_hash in r.tombstoned:
self._in_flight.discard(mm_hash)
self._cache.discard(mm_hash)
self._tombstones.add(mm_hash)
logger.debug("EC consumer: NIXL xfer failed mm_hash=%s", mm_hash)
for mm_hash in r.quarantined:
# DMA still running: keep the blocks reserved (entry stays
# not-ready, hence non-evictable) until the xfer settles.
self._in_flight.discard(mm_hash)
self._tombstones.add(mm_hash)
logger.debug(
"EC consumer: NIXL xfer mm_hash=%s quarantined; DMA still running",
mm_hash,
)
for mm_hash in r.retryable:
# Release this attempt's blocks, retaining the request wait budget.
self._in_flight.discard(mm_hash)
self._cache.discard(mm_hash)
logger.debug(
"EC consumer: NIXL xfer mm_hash=%s retryable; will re-request",
mm_hash,
)
for mm_hash, _block_indices in r.settled:
self._cache.discard(mm_hash)
logger.debug(
"EC consumer: quarantined NIXL xfer mm_hash=%s settled", mm_hash
)
def _on_peer_down(self, addr: Any) -> None:
session = self._sessions.pop(addr, None)
if session is None:
return
session.on_peer_down()
self._process_session_results(session)
session.close()
logger.info("EC consumer: producer peer down addr=%s", addr)
def _promote_completed_reads(self) -> None:
"""Drain the just-completed set built during ``_poll_step``.
Reads are marked ready in ``_process_session_results``, so a deferred
request is re-admitted by ``ensure_cache_available`` and loaded through
the same local path as a natively cached entry
(``update_state_after_alloc`` -> ``_pending_loads``). No explicit
pin/load bookkeeping is needed here beyond clearing the set.
"""
self._step_completed.clear()
def _num_blocks(self, shape: tuple[int, ...]) -> int:
size_bytes = shape[0] * shape[1] * self._dtype.itemsize
block_size = self._region.block_size_bytes
return (size_bytes + block_size - 1) // block_size
def update_state_after_alloc(self, request: "Request", index: int) -> None:
feature = request.mm_features[index]
mm_hash = feature.identifier
if self._is_producer and self._cache.get(mm_hash) is None:
shape = (
feature.mm_position.get_num_embeds(),
get_encoder_output_width(self._vllm_config, feature.modality),
)
entry = self._cache.alloc(mm_hash, self._num_blocks(shape), shape)
if entry is not None:
self._pending_saves[mm_hash] = list(entry.block_ids)
if self._is_consumer and mm_hash not in self._pending_loads:
entry = self._cache.get(mm_hash)
if entry is not None and entry.ready:
assert entry.shape is not None
self._cache.pin(mm_hash)
transfer_id = self._next_transfer_id
self._next_transfer_id += 1
self._pending_loads[mm_hash] = (
transfer_id,
list(entry.block_ids),
entry.shape,
)
self._load_acks[transfer_id] = (mm_hash, self._expected_load_reports)
def build_connector_meta(
self, scheduler_output: "SchedulerOutput"
) -> ECCPUConnectorMetadata:
meta = ECCPUConnectorMetadata()
if self._is_producer:
if self._nixl_enabled and self._producer_session is not None:
self._producer_session.poll_step()
meta.saves = self._pending_saves
self._pending_saves = {}
if self._is_consumer:
if self._nixl_enabled:
self._promote_completed_reads()
meta.loads = self._pending_loads
self._pending_loads = {}
if self._nixl_enabled:
self._first_in_batch = True
return meta
def update_connector_output(self, connector_output: "ECConnectorOutput") -> None:
"""Apply the worker's memcpy-completion report to the cache.
Completed saves become safe to mark ready. A load is unpinned once
every participating rank has reported its transfer id; reports for a
transfer that has already been released, or for one this scheduler
never dispatched, are ignored.
"""
meta = connector_output.ec_connector_worker_meta
if not isinstance(meta, ECCPUWorkerMetadata):
return
for mm_hash in meta.completed_saves:
entry = self._cache.get(mm_hash)
if entry is None:
logger.debug(
"EC producer: worker reported completed save for unknown "
"mm_hash=%s (already discarded/evicted?)",
mm_hash,
)
elif not entry.ready:
self._cache.mark_ready(mm_hash)
logger.debug("EC producer: mm_hash=%s marked ready", mm_hash)
for transfer_id in meta.completed_loads:
pending = self._load_acks.get(transfer_id)
if pending is None:
continue
mm_hash, outstanding = pending
if outstanding > 1:
self._load_acks[transfer_id] = (mm_hash, outstanding - 1)
continue
# Drop the entry before unpinning so a replayed report is treated
# as stale rather than releasing the pin a second time.
del self._load_acks[transfer_id]
self._cache.unpin(mm_hash)
logger.debug("EC consumer: mm_hash=%s unpinned after load", mm_hash)
def has_pending_push_work(self) -> bool:
"""Keep the engine stepping so this connector's polls keep running.
With NIXL enabled the engine tick is the only driver of
``ProducerSession.poll_step()``, and a producer's work arrives from a
remote consumer, so nothing local would otherwise wake it. Without NIXL
this only has to outlive dispatched saves and loads, so the engine can
quiesce once the worker has confirmed them.
"""
return self._nixl_enabled or self._cache.has_held_entries()
def request_finished(
self, request: "Request"
) -> tuple[bool, "dict[str, Any] | None"]:
for feature in request.mm_features:
self._deferred_since.pop((request.request_id, feature.identifier), None)
if not (self._nixl_enabled and self._is_producer):
return False, None
if not request.mm_features:
return False, None
items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver)
for feature in request.mm_features:
mm_hash = feature.identifier
entry = self._cache.get(mm_hash)
if entry is None:
# Never saved, or evicted since. Publishing placeholder
# metadata without an address invites an orchestrator to
# rewrite the media into a reference to an encoding no
# consumer can fetch, leaving the decoder nothing to embed.
# Publish neither, so the media stays on the request.
items[mm_hash]["metadata"] = {}
logger.debug(
"EC producer: mm_hash=%s absent at request_finished; "
"announcing no metadata so the media is not rewritten away",
mm_hash,
)
continue
# Announce even if the save's GPU->mmap copy hasn't been
# confirmed complete yet: a not-ready entry can't be evicted, so
# it will still be here by the time a consumer's XferReq arrives.
# A read arriving before the save lands is NACKed NACK_NOT_READY,
# which the consumer retries on a later step rather than treating
# as a miss.
assert entry.shape is not None
size_bytes = len(entry.block_ids) * self._region.block_size_bytes
items[mm_hash].update(
peer_host=self._peer_host,
peer_port=self._peer_port,
size_bytes=size_bytes,
shape=list(entry.shape),
)
logger.debug(
"EC producer: announcing NIXL-readable encodings req_id=%s items=%s",
request.request_id,
items,
)
return False, items
def shutdown(self) -> None:
self._pending_saves.clear()
self._pending_loads.clear()
self._load_acks.clear()
self._deferred_since.clear()
self._is_producer = False
self._is_consumer = False
if self._nixl_enabled:
self._teardown_nixl()
try:
self._region.cleanup()
except Exception:
logger.debug("ec: region cleanup failed", exc_info=True)
def _teardown_nixl(self) -> None:
"""Close the control sockets and release the NIXL agent.
Guards each step and clears what it releases, so it is safe on a
scheduler whose NIXL setup failed part-way and safe to call twice.
The shared region is not touched — shutdown() owns that.
"""
if self._producer_session is not None:
try:
self._producer_session.close()
except Exception:
logger.debug("ec: producer session close failed", exc_info=True)
self._producer_session = None
for session in list(self._sessions.values()):
try:
session.close()
except Exception:
logger.debug("ec: consumer session close failed", exc_info=True)
self._sessions.clear()
if self._transport is not None:
try:
self._transport.close()
except Exception:
logger.debug("ec: client transport close failed", exc_info=True)
self._transport = None
if self._data is not None:
try:
self._data.deregister()
except Exception:
logger.debug("ec: deregister failed", exc_info=True)
self._data = None
vllm/distributed/ec_transfer/ec_connector/cpu/scheduler/embedding_cache.py
-184
View File
@@ -1,184 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""EmbeddingCache — named-entry block cache with FIFO eviction.
Manages a fixed pool of block IDs keyed by content identity (mm_hash).
Entries transition through: not-ready → ready (evictable) → pinned.
Eviction targets ready + unpinned entries in FIFO order.
All public methods are thread-safe.
"""
import threading
from collections import OrderedDict
class CacheEntry:
"""A single cache entry. Read `.block_ids`, `.shape` and `.ready` freely;
mutations only through EmbeddingCache methods.
`shape` is the encoder output the blocks hold, so a load can restore it
without knowing how wide this model's encoder output is.
"""
__slots__ = ("block_ids", "shape", "_pin_count")
def __init__(
self, block_ids: tuple[int, ...], shape: tuple[int, int] | None = None
) -> None:
self.block_ids = block_ids
self.shape = shape
self._pin_count = -1 # not ready
@property
def ready(self) -> bool:
return self._pin_count >= 0
@property
def evictable(self) -> bool:
return self._pin_count == 0
def mark_ready(self):
self._pin_count = 0
def pin(self):
self._pin_count += 1
class EmbeddingCache:
"""Fixed-size block cache with FIFO eviction.
Entries are keyed by content identity (e.g. mm_hash). Blocks are
allocated from a free-set; when space is needed, ready + unpinned
entries are evicted oldest-first.
The caller is responsible for deciding *when* to call ``mark_ready``
(e.g. after enough engine steps have elapsed for the worker to have
completed the write).
"""
def __init__(self, num_blocks: int) -> None:
self._num_blocks = num_blocks
self._lock = threading.Lock()
# Available block IDs, LIFO stack.
self._free_blocks: list[int] = list(range(num_blocks))
# Key → live entry (all states: not-ready, ready, pinned).
self._entries: dict[str, CacheEntry] = {}
# Ready + unpinned entries in insertion order (FIFO eviction).
self._entries_free_list: OrderedDict[str, None] = OrderedDict()
# Total blocks reclaimable via eviction (fail-fast check).
self._evictable_block_count: int = 0
def get(self, key: str) -> CacheEntry | None:
"""Return the entry for *key*, or None if not present."""
with self._lock:
return self._entries.get(key)
def has_held_entries(self) -> bool:
"""True if any entry is not-ready or pinned (an in-flight transfer).
Held entries are exactly those absent from the eviction free list, so
they represent saves awaiting mark_ready or loads awaiting unpin.
"""
with self._lock:
return len(self._entries) != len(self._entries_free_list)
def alloc(
self, key: str, n_blocks: int, shape: tuple[int, int] | None = None
) -> CacheEntry | None:
"""Allocate *n_blocks* for *key*, evicting as needed.
The entry starts not-ready. Returns None if there is not enough
space even after evicting all evictable entries.
"""
with self._lock:
assert key not in self._entries, (
f"EmbeddingCache: duplicate alloc for {key!r}"
)
assert n_blocks <= self._num_blocks, (
f"EmbeddingCache: {n_blocks} blocks requested but capacity "
f"is {self._num_blocks}"
)
if len(self._free_blocks) + self._evictable_block_count < n_blocks:
return None
self._evict_until(n_blocks)
# Canonical ascending order: keeps a freed entry's blocks contiguous
# in the free list when it is evicted, and lets consumers detect
# which of an entry's blocks are adjacent.
block_ids = tuple(sorted(self._free_blocks.pop() for _ in range(n_blocks)))
entry = CacheEntry(block_ids, shape)
self._entries[key] = entry
return entry
def mark_ready(self, key: str) -> None:
"""Mark an entry as ready (data is CPU-visible)."""
with self._lock:
entry = self._entries[key]
assert entry._pin_count == -1, (
f"EmbeddingCache: mark_ready on already-ready entry {key!r}"
)
entry.mark_ready()
self._entries_free_list[key] = None
self._evictable_block_count += len(entry.block_ids)
def pin(self, key: str) -> None:
"""Pin an entry (prevent eviction)."""
with self._lock:
entry = self._entries[key]
assert entry.ready, f"EmbeddingCache: pin of not-ready entry {key!r}"
if entry.evictable:
del self._entries_free_list[key]
self._evictable_block_count -= len(entry.block_ids)
entry.pin()
def pin_if_ready(self, key: str) -> CacheEntry | None:
"""Atomically pin *key* if present and ready; return its entry.
Returns None when the key is absent. An entry returned with `ready`
False was not pinned: its save is still in flight, which the producer
reports differently from a miss because such a read can be retried.
"""
with self._lock:
entry = self._entries.get(key)
if entry is None or not entry.ready:
return entry
if entry.evictable:
del self._entries_free_list[key]
self._evictable_block_count -= len(entry.block_ids)
entry.pin()
return entry
def unpin(self, key: str) -> None:
"""Unpin an entry. Asserts currently pinned."""
with self._lock:
entry = self._entries[key]
assert entry._pin_count > 0, (
f"EmbeddingCache: unpin of unpinned entry {key!r}"
)
entry._pin_count -= 1
if entry._pin_count == 0:
self._entries_free_list[key] = None
self._evictable_block_count += len(entry.block_ids)
def discard(self, key: str) -> None:
"""Remove a not-ready in-flight entry, returning its blocks to the pool.
Used when an in-flight fill (e.g. a NIXL READ) fails before the entry
is marked ready. Asserts the entry is present and not ready — ready or
pinned entries are reclaimed through eviction/unpin, not discard.
"""
with self._lock:
entry = self._entries[key]
assert entry._pin_count == -1, (
f"EmbeddingCache: discard of ready/pinned entry {key!r}"
)
del self._entries[key]
self._free_blocks.extend(entry.block_ids)
def _evict_until(self, n_blocks: int) -> None:
"""Evict ready+unpinned entries FIFO until enough space. Lock held."""
while len(self._free_blocks) < n_blocks:
key, _ = self._entries_free_list.popitem(last=False)
entry = self._entries.pop(key)
self._evictable_block_count -= len(entry.block_ids)
self._free_blocks.extend(entry.block_ids)
vllm/distributed/ec_transfer/ec_connector/cpu/session.py
-644
View File
@@ -1,644 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Session objects for the ECCPUConnector.
Layer hierarchy: connector → session → transport (connection)
Two long-lived sessions:
ProducerSession — owns ZmqServerTransport. The connector calls poll_step()
once per engine step, which drains the transport without blocking and
feeds the raw bytes to poll(). poll() decodes XferReqs, grants or NACKs
them, pins source blocks, sends XferAck replies, drains NIXL completion
notifications, and sweeps expired xfers.
ConsumerSession — owns one ZmqClientConnection and the NIXL agent
registration for that peer. The connector calls poll(messages, now)
once per engine step with the raw bytes collected for this peer.
poll() decodes XferAcks, dispatches them to ConsumerXfer objects,
and advances all active xfer state machines.
Internal data classes:
ProducerXfer — one pinned read grant, keyed by session_id.
ConsumerXfer — state machine for one NIXL READ, keyed by mm_hash.
"""
import enum
import logging
import time
import uuid
from dataclasses import dataclass, field
from typing import Any
import msgspec
from vllm.distributed.ec_transfer.ec_connector.cpu.control.zmq import (
ZmqClientConnection,
ZmqClientTransport,
ZmqServerTransport,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.data.base import (
DataTransport,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.protocol import (
EC_CONNECTOR_VERSION,
EXPECTED_NACKS,
RETRYABLE_NACKS,
XferAck,
XferReq,
XferStatus,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.scheduler.embedding_cache import (
EmbeddingCache,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.utils import PeerAddr
from vllm.logger import init_logger
logger = init_logger(__name__)
PRODUCER_PIN_LEASE_S = 30.0
_CONSUMER_READ_TIMEOUT_S = 20.0
_CONSUMER_QUARANTINE_TIMEOUT_S = 60.0
# ── XferState ─────────────────────────────────────────────────────────────────
class XferState(enum.Enum):
"""Lifecycle states of a single ConsumerXfer."""
WAITING_ACK = "waiting_ack" # XferReq sent, awaiting XferAck
READING = "reading" # NIXL READ in flight
DONE = "done" # READ completed; promote to local cache
ACK_TIMEOUT = "ack_timeout" # XferAck never arrived; free blocks + retry
READ_FAILED = "read_failed" # Unexpected NIXL state; free blocks + tombstone
QUARANTINED = "quarantined" # Read timed out, NIXL unabortable; keep blocks
SETTLED = "settled" # Quarantined DMA done; safe to free blocks
# ── ProducerXfer ──────────────────────────────────────────────────────────────
@dataclass
class ProducerXfer:
"""One pinned read grant on the producer side.
Pure data — lifecycle operations (pin/unpin) are performed by ProducerSession.
Keyed by (consumer_session_id, mm_hash) in ProducerSession._active_xfers.
"""
mm_hash: str
block_indices: list[int]
deadline: float
def is_expired(self) -> bool:
return time.monotonic() > self.deadline
# ── ConsumerXfer ──────────────────────────────────────────────────────────────
@dataclass
class ConsumerXfer:
"""State machine for one NIXL READ on the consumer side.
transfer_handle is the in-flight NIXL handle returned by post_read() —
used for check_xfer_state() / release_xfer_handle(). The dlist handle
for the remote peer lives inside DataTransport, not here.
"""
mm_hash: str
block_indices: list[int]
addr: PeerAddr
deadline: float
data: DataTransport
consumer_session_id: str # session identity; combined with mm_hash for notif_msg
transfer_handle: Any | None = field(default=None)
_quarantined: bool = field(default=False, init=False)
def handle_ack(self, ack: XferAck, agent_name: str) -> bool:
"""Issue the NIXL READ on an OK XferAck. Returns False on any NACK."""
if ack.status != XferStatus.OK:
return False
# notif_msg encodes both session_id and mm_hash so the producer can
# do a direct (session_id, mm_hash) → ProducerXfer lookup on completion.
notif_msg = f"{self.consumer_session_id}:{self.mm_hash}".encode()
self.transfer_handle = self.data.post_read(
self.block_indices,
agent_name,
ack.src_block_indices,
notif_msg=notif_msg,
)
self.deadline = time.monotonic() + _CONSUMER_READ_TIMEOUT_S
logger.debug(
"EC consumer: NIXL READ posted mm_hash=%s agent=%s local_blocks=%d "
"remote_blocks=%d",
self.mm_hash,
agent_name,
len(self.block_indices),
len(ack.src_block_indices),
)
return True
def poll(self, now: float) -> XferState:
if self._quarantined:
return self._poll_quarantined(now)
if self.transfer_handle is None:
return (
XferState.ACK_TIMEOUT if now > self.deadline else XferState.WAITING_ACK
)
return self._poll_reading(now)
def _poll_reading(self, now: float) -> XferState:
assert self.transfer_handle is not None
try:
state = self.data.check_xfer_state(self.transfer_handle)
except Exception:
logger.exception(
"EC Connector: check_xfer_state failed for mm_hash=%s", self.mm_hash
)
self.data.release_xfer_handle(self.transfer_handle)
self.transfer_handle = None
return XferState.READ_FAILED
if state == "DONE":
self.data.release_xfer_handle(self.transfer_handle)
self.transfer_handle = None
logger.debug("EC consumer: NIXL READ completed mm_hash=%s", self.mm_hash)
return XferState.DONE
if state == "PROC":
if now > self.deadline:
self._quarantined = True
self.deadline = now + _CONSUMER_QUARANTINE_TIMEOUT_S
logger.warning(
"EC Connector: READ for mm_hash=%s from %s:%d timed out; "
"quarantining",
self.mm_hash,
self.addr[0],
self.addr[1],
)
return XferState.QUARANTINED
return XferState.READING
logger.warning(
"EC Connector: READ for mm_hash=%s unexpected NIXL state %r",
self.mm_hash,
state,
)
self.data.release_xfer_handle(self.transfer_handle)
self.transfer_handle = None
return XferState.READ_FAILED
def _poll_quarantined(self, now: float) -> XferState:
if self.transfer_handle is None:
return XferState.SETTLED
# If quarantine has been too long, assume NIXL is stuck and evict blocks.
if now > self.deadline:
self.data.release_xfer_handle(self.transfer_handle)
self.transfer_handle = None
logger.warning(
"EC Connector: quarantined mm_hash=%s from %s:%d timed out; "
"evicting blocks",
self.mm_hash,
self.addr[0],
self.addr[1],
)
return XferState.READ_FAILED
try:
state = self.data.check_xfer_state(self.transfer_handle)
except Exception:
logger.exception(
"EC Connector: check_xfer_state failed for quarantined mm_hash=%s",
self.mm_hash,
)
state = None
if state == "PROC":
return XferState.QUARANTINED
self.data.release_xfer_handle(self.transfer_handle)
self.transfer_handle = None
logger.debug(
"EC Connector: quarantined mm_hash=%s settled (state=%s)",
self.mm_hash,
state,
)
return XferState.SETTLED
def cancel(self) -> None:
"""Cancel a WAITING_ACK xfer (peer down). No tombstone — caller allows retry."""
assert self.transfer_handle is None, "cancel() requires WAITING_ACK state"
def release(self) -> None:
"""Force-release the transfer handle. Shutdown only."""
if self.transfer_handle is not None:
self.data.release_xfer_handle(self.transfer_handle)
self.transfer_handle = None
# ── ProducerSession ───────────────────────────────────────────────────────────
class ProducerSession:
"""Long-lived producer session. Owns ZmqServerTransport.
poll_step() runs once per engine step: drains the transport without
blocking, decodes XferReqs, grants or NACKs each one, drains NIXL completion
notifications, and sweeps expired xfers.
"""
def __init__(
self,
transport: ZmqServerTransport,
data: DataTransport,
cache: EmbeddingCache,
compat_hash: str,
) -> None:
self._transport = transport
self._data = data
self._cache = cache # EmbeddingCache — single block owner
self._compat_hash = compat_hash
# Key: "{consumer_session_id}:{mm_hash}" — matches the NIXL notif_msg exactly,
# so completion notification lookup is a direct dict pop with no parsing.
self._active_xfers: dict[str, ProducerXfer] = {}
self._req_decoder = msgspec.msgpack.Decoder(XferReq)
self._req_encoder = msgspec.msgpack.Encoder()
def close(self) -> None:
for xfer in self._active_xfers.values():
self._cache.unpin(xfer.mm_hash) # release any still-held pins
self._active_xfers.clear()
self._transport.close()
def poll_step(self) -> None:
"""Drain the transport without blocking and process what arrived."""
self.poll(self._transport.poll(timeout_ms=0))
def poll(self, messages: list[tuple[bytes, bytes]]) -> None:
"""Decode and process inbound XferReqs, drain NIXL notifs, sweep timeouts.
messages: (identity, payload) pairs from ZmqServerTransport.poll().
"""
for identity, payload in messages:
try:
req = self._req_decoder.decode(payload)
except (msgspec.DecodeError, msgspec.ValidationError):
logger.warning("EC Connector: dropped malformed XferReq")
continue
ack = self._grant_or_nack(req)
self._transport.send(identity, self._req_encoder.encode(ack))
self._drain_notifs()
self._sweep_timeouts()
def _grant_or_nack(self, req: XferReq) -> XferAck:
if req.connector_version != EC_CONNECTOR_VERSION:
logger.warning(
"EC Connector: incompatible version for mm_hash=%s", req.mm_hash
)
return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_VERSION)
if req.compatibility_hash != self._compat_hash:
logger.warning(
"EC Connector: incompatible compat hash for mm_hash=%s", req.mm_hash
)
return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INCOMPAT)
if not req.session_id:
# Grants are keyed by (session_id, mm_hash), so an omitted id from
# one peer aliases another's: the first completion would unpin
# blocks the other peer is still reading.
logger.warning(
"EC Connector: XferReq without session_id for mm_hash=%s", req.mm_hash
)
return XferAck(mm_hash=req.mm_hash, status=XferStatus.NACK_INTERNAL)
entry = self._cache.pin_if_ready(req.mm_hash)
if entry is None or not entry.ready:
status = (
XferStatus.NACK_MISSING if entry is None else XferStatus.NACK_NOT_READY
)
logger.log(
logging.INFO if status in EXPECTED_NACKS else logging.WARNING,
"EC Connector: mm_hash=%s not serveable in local cache (%s); NACKing",
req.mm_hash,
status.name,
)
return XferAck(mm_hash=req.mm_hash, status=status)
block_indices = list(entry.block_ids)
key = f"{req.session_id}:{req.mm_hash}"
if key in self._active_xfers:
# This peer is re-requesting a grant it still holds, having timed
# out waiting. Overwriting it drops the pin taken for it, which
# neither a completion nor the sweep would ever release.
self._cache.unpin(req.mm_hash)
self._active_xfers[key] = ProducerXfer(
mm_hash=req.mm_hash,
block_indices=block_indices,
deadline=time.monotonic() + PRODUCER_PIN_LEASE_S,
)
logger.debug(
"EC producer: serving NIXL READ request mm_hash=%s key=%s blocks=%d",
req.mm_hash,
key,
len(block_indices),
)
return XferAck(
mm_hash=req.mm_hash,
status=XferStatus.OK,
session_id=req.session_id, # echo consumer's session_id back
src_block_indices=block_indices,
agent_metadata=self._data.get_agent_metadata(),
mem_descriptor=self._data.get_mem_descriptor(),
)
def _drain_notifs(self) -> None:
try:
notifs = self._data.get_new_notifs()
except Exception:
logger.exception("EC Connector: get_new_notifs failed")
return
for msgs in notifs.values():
for msg in msgs:
try:
key = msg.decode("utf-8") # "{session_id}:{mm_hash}"
except Exception:
logger.exception("EC Connector: Dropped malformed NIXL notif")
continue
xfer = self._active_xfers.pop(key, None)
if xfer is not None:
self._cache.unpin(xfer.mm_hash)
logger.debug(
"EC producer: NIXL READ completed mm_hash=%s key=%s",
xfer.mm_hash,
key,
)
def _sweep_timeouts(self) -> None:
for key, xfer in list(self._active_xfers.items()):
if xfer.is_expired():
logger.warning(
"EC producer: consumer did not complete NIXL READ "
"mm_hash=%s key=%s within the pin lease; releasing "
"pinned blocks",
xfer.mm_hash,
key,
)
self._cache.unpin(xfer.mm_hash)
del self._active_xfers[key]
# ── ConsumerSession ───────────────────────────────────────────────────────────
@dataclass
class ConsumerSessionResults:
"""Results from ConsumerSession.take_results()."""
completed: set[str] # promote to local cache
tombstoned: set[str] # free blocks + prevent retry this step
quarantined: set[str] # timed-out, DMA still running: prevent retry
# but DO NOT free blocks — session holds them
retryable: set[str] # free blocks, allow another attempt (not ready, peer down)
settled: list[tuple[str, list[int]]] # (mm_hash, block_indices) quarantine cleared
class ConsumerSession:
"""Long-lived per-producer session on the consumer side.
Owns the ZMQ DEALER connection and the NIXL agent registration for one
producer peer. The connector calls poll(messages, now) once per engine
step with the raw bytes collected from this peer's DEALER socket.
poll() decodes XferAcks, dispatches them to the right ConsumerXfer,
advances all active xfer state machines, and drains quarantined xfers.
Results accumulate in internal sets; call take_results() to collect them.
"""
def __init__(
self,
addr: PeerAddr,
zmq_conn: ZmqClientConnection,
transport: ZmqClientTransport,
data: DataTransport,
compat_hash: str,
) -> None:
self._addr = addr
self._zmq = zmq_conn
self._transport = transport
self._data = data
self._compat_hash = compat_hash
self._session_id = str(
uuid.uuid4()
) # stable identity for (session_id:mm_hash) keys
self._nixl_agent_name: str | None = None
self._nixl_metadata_bytes: bytes | None = None
self._xfers: dict[str, ConsumerXfer] = {} # mm_hash → xfer
self._quarantined: list[ConsumerXfer] = [] # DMA pending, to poll
self._newly_quarantined: set[str] = set() # quarantined since last take_results
self._completed: set[str] = set()
self._tombstoned: set[str] = set()
self._retryable: set[str] = set()
self._settled: list[tuple[str, list[int]]] = []
self._encoder = msgspec.msgpack.Encoder()
self._decoder = msgspec.msgpack.Decoder(XferAck)
@property
def addr(self) -> PeerAddr:
return self._addr
def start_xfer(
self, mm_hash: str, block_indices: list[int], deadline: float
) -> None:
"""Create a ConsumerXfer and send the XferReq over the connection."""
xfer = ConsumerXfer(
mm_hash=mm_hash,
block_indices=block_indices,
addr=self._addr,
deadline=deadline,
data=self._data,
consumer_session_id=self._session_id,
)
self._zmq.send(
self._encoder.encode(
XferReq(
mm_hash=mm_hash,
compatibility_hash=self._compat_hash,
session_id=self._session_id,
)
)
)
self._xfers[mm_hash] = xfer
logger.debug(
"EC Connector: XferReq sent mm_hash=%s to %s:%d",
mm_hash,
self._addr[0],
self._addr[1],
)
def poll(self, messages: list[bytes], now: float) -> None:
"""Decode inbound XferAcks, dispatch to xfers, and advance all state machines.
messages: raw bytes from ZmqClientTransport.poll() for this peer.
Results accumulate until take_results() is called.
"""
for raw in messages:
try:
ack = self._decoder.decode(raw)
except (msgspec.DecodeError, msgspec.ValidationError):
logger.warning(
"EC Connector: dropped malformed XferAck from %s:%d",
self._addr[0],
self._addr[1],
)
continue
self._dispatch_ack(ack)
for mm_hash, xfer in list(self._xfers.items()):
state = xfer.poll(now)
if state == XferState.DONE:
self._completed.add(mm_hash)
del self._xfers[mm_hash]
elif state == XferState.ACK_TIMEOUT:
# The scheduler bounds retries with a per-request deadline.
logger.warning(
"EC consumer: no XferAck for mm_hash=%s from %s:%d in time; "
"will re-request",
mm_hash,
self._addr[0],
self._addr[1],
)
self._retryable.add(mm_hash)
del self._xfers[mm_hash]
elif state == XferState.READ_FAILED:
self._tombstoned.add(mm_hash)
del self._xfers[mm_hash]
elif state == XferState.QUARANTINED:
# DMA may still be writing — caller must NOT free blocks.
self._quarantined.append(xfer)
self._newly_quarantined.add(mm_hash)
del self._xfers[mm_hash]
self._drain_quarantine(now)
def _dispatch_ack(self, ack: XferAck) -> None:
xfer = self._xfers.get(ack.mm_hash)
if xfer is None or xfer.transfer_handle is not None:
return # unknown mm_hash or duplicate ack
if ack.status != XferStatus.OK:
logger.log(
logging.INFO if ack.status in EXPECTED_NACKS else logging.WARNING,
"EC Connector: NACK %s from %s:%d for mm_hash=%s",
ack.status.name,
self._addr[0],
self._addr[1],
ack.mm_hash,
)
del self._xfers[ack.mm_hash]
# The producer announced this encoding and its save is still in
# flight, so another attempt on a later step can be granted. Every
# other NACK is terminal for this mm_hash.
if ack.status in RETRYABLE_NACKS:
self._retryable.add(ack.mm_hash)
else:
self._tombstoned.add(ack.mm_hash)
return
try:
agent_name = self._ensure_registered(ack.agent_metadata, ack.mem_descriptor)
except Exception:
logger.exception(
"EC Connector: failed to register peer %s:%d for mm_hash=%s",
self._addr[0],
self._addr[1],
ack.mm_hash,
)
del self._xfers[ack.mm_hash]
self._tombstoned.add(ack.mm_hash)
return
if not xfer.handle_ack(ack, agent_name):
del self._xfers[ack.mm_hash]
self._tombstoned.add(ack.mm_hash)
def _ensure_registered(self, metadata: bytes, mem_descriptor: bytes) -> str:
"""Register or re-register the peer in DataTransport. Returns agent_name."""
if self._nixl_agent_name is not None and self._nixl_metadata_bytes == metadata:
return self._nixl_agent_name
if self._nixl_agent_name is not None:
logger.info(
"EC Connector: producer %s:%d restarted; re-registering",
self._addr[0],
self._addr[1],
)
self._data.remove_remote_peer(self._nixl_agent_name)
agent_name = self._data.add_remote_peer(metadata, mem_descriptor)
self._nixl_agent_name = agent_name
self._nixl_metadata_bytes = metadata
logger.debug(
"EC Connector: registered peer %s:%d agent=%s",
self._addr[0],
self._addr[1],
agent_name,
)
return agent_name
def _drain_quarantine(self, now: float) -> None:
still_pending: list[ConsumerXfer] = []
for xfer in list(self._quarantined):
state = xfer.poll(now)
if state == XferState.SETTLED:
self._settled.append((xfer.mm_hash, xfer.block_indices))
else:
still_pending.append(xfer)
self._quarantined = still_pending
def on_peer_down(self) -> None:
"""WAITING_ACK xfers → cancel (retry allowed); READING xfers → quarantine."""
for mm_hash, xfer in list(self._xfers.items()):
if xfer.transfer_handle is None:
xfer.cancel()
self._retryable.add(mm_hash)
else:
self._quarantined.append(xfer)
self._newly_quarantined.add(mm_hash)
self._xfers.clear()
def take_results(self) -> ConsumerSessionResults:
"""Return and clear all accumulated results since the last call."""
results = ConsumerSessionResults(
completed=self._completed,
tombstoned=self._tombstoned,
quarantined=self._newly_quarantined,
retryable=self._retryable,
settled=self._settled,
)
self._completed = set()
self._tombstoned = set()
self._newly_quarantined = set()
self._retryable = set()
self._settled = []
return results
def has_in_flight(self) -> bool:
return bool(self._xfers) or bool(self._quarantined)
def close(self) -> None:
"""Release all transfer handles, deregister NIXL peer, close connection.
Connection teardown is routed through the transport so the pooled
connection is dropped in lock-step with the session — a session and its
DEALER connection have a 1:1 lifecycle, so a later connect() to the same
peer builds a fresh socket instead of handing back this dead one.
"""
for xfer in self._xfers.values():
xfer.release()
self._xfers.clear()
for xfer in self._quarantined:
xfer.release()
self._quarantined.clear()
if self._nixl_agent_name is not None:
self._data.remove_remote_peer(self._nixl_agent_name)
self._transport.remove(self._addr)
vllm/distributed/ec_transfer/ec_connector/cpu/utils.py
-47
View File
@@ -1,47 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Supporting utilities for the ECCPUConnector scheduler."""
import msgspec
from vllm.logger import init_logger
# (host, port) pair identifying a remote peer's ZMQ side-channel endpoint.
PeerAddr = tuple[str, int]
logger = init_logger(__name__)
# Msgpack (de)serialization for the list-of-(addr, size, device_id)
# descriptor form that NIXL's `get_xfer_descs` expects.
_MemDescList = list[tuple[int, int, int]]
_mem_desc_encoder = msgspec.msgpack.Encoder()
_mem_desc_decoder = msgspec.msgpack.Decoder(_MemDescList)
def serialize_mem_descriptor(descs: _MemDescList) -> bytes:
return _mem_desc_encoder.encode(descs)
def deserialize_mem_descriptor(payload: bytes) -> _MemDescList:
return _mem_desc_decoder.decode(payload)
def build_block_descs(
base_ptr: int,
num_blocks: int,
block_size_bytes: int,
device_id: int = 0,
) -> _MemDescList:
"""Per-block `(addr, size, device_id)` tuples for an `ECSharedRegion`.
Handed to `nixl.get_xfer_descs` to build a dlist where descriptor `i`
addresses block `i` in the mmap. This makes a WRITE to block indices
`[i1, i2, ...]` a simple `make_prepped_xfer(..., [i1, i2])` on both
ends — no per-request address math.
"""
return [
(base_ptr + i * block_size_bytes, block_size_bytes, device_id)
for i in range(num_blocks)
]
vllm/distributed/ec_transfer/ec_connector/cpu/worker/__init__.py
-455
View File
@@ -1,455 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Worker-side of the ECCPUConnector.
Thin, stateless across steps: opens the shared mmap region and uses the
per-step connector metadata (`ECCPUConnectorMetadata`) to decide which
blocks to copy in each direction.
"""
from collections import deque
from collections.abc import Sequence
from dataclasses import dataclass
from enum import Enum, auto
from itertools import chain
from typing import TYPE_CHECKING
import numpy as np
import torch
from vllm._custom_ops import swap_blocks_batch
from vllm.distributed.ec_transfer.ec_connector.cpu.common import (
ECCPUConnectorMetadata,
ECCPUWorkerMetadata,
create_ec_shared_region,
)
from vllm.distributed.ec_transfer.ec_connector.cpu.worker.descriptor_buffers import (
DescriptorBufferPool,
DescriptorBuffers,
)
from vllm.distributed.parallel_state import (
get_pcp_group,
get_tensor_model_parallel_rank,
)
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.platform_utils import is_pin_memory_available
if TYPE_CHECKING:
from vllm.config import VllmConfig
logger = init_logger(__name__)
class ECCPUTransferDirection(Enum):
HOST_TO_DEVICE = auto()
DEVICE_TO_HOST = auto()
def _coalesce_runs(
block_ids: Sequence[int],
) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Split block ids into runs that are consecutive in the region.
Returns `(slots, first_blocks, num_blocks)`, one element per run, where
`slots` is the run's offset in the flat block sequence. Per-descriptor
overhead dominates these copies, so describing a run with one descriptor
rather than one per block is worth an order of magnitude.
Both callers derive the non-region side of a run from its slot, so that
side advances in lockstep with the region and only the region has to be
split. `EmbeddingCache` hands out ascending ids, which is what makes runs
findable at all.
"""
ids = np.fromiter(block_ids, dtype=np.int64, count=len(block_ids))
breaks = np.flatnonzero(np.diff(ids) != 1) + 1
slots = np.concatenate(([0], breaks))
ends = np.concatenate((breaks, [ids.size]))
return slots, ids[slots], ends - slots
@dataclass
class Transfer:
"""A batched copy in flight on a dedicated GPU stream.
Its descriptor buffers and stream are held until `end_event` fires; only
then are the buffers and stream safe to recycle and the `completions` safe
to report. `start_event`/`end_event` bracket the copy so its elapsed time
can be measured on completion.
`completions` is what the scheduler is told about this copy: mm_hashes for
saves, transfer ids for loads.
"""
start_event: torch.Event
end_event: torch.Event
completions: list[str] | list[int]
bufs: DescriptorBuffers
stream: torch.Stream
num_bytes: int
class ECCPUWorker:
"""Worker-side delegate for the ECCPUConnector.
- Producer role: copies `encoder_cache[mm_hash]` → `mmap[block_ids]`
for each entry in `metadata.saves`. Descriptor buffers are filled
directly in `save_caches`; the actual DMA is issued as a single
batched call in `flush_saves`.
- Consumer role: copies `mmap[block_ids]` → `encoder_cache[mm_hash]`
for all entries in `metadata.loads` via a single `swap_blocks_batch`
call on a dedicated stream.
- On `ec_both` nodes both paths run back-to-back in a single step.
Every batched copy runs on a stream drawn from a pool so it overlaps the
compute stream, bracketed by a start/end event pair. The stream, events,
and descriptor buffers are recycled once the end event fires.
Both paths move memory between the two streams, so each buffer is registered
against the stream that is not its own: save sources against the save stream,
load destinations against the compute stream. Without that, the caching
allocator would be free to reissue memory a copy or a queued read still
needs, since it only tracks the stream a buffer was allocated on.
"""
def __init__(self, vllm_config: "VllmConfig") -> None:
self._region = create_ec_shared_region(vllm_config)
# Model dtype; used to reinterpret raw int8 blocks on load.
self._dtype = vllm_config.model_config.dtype
self._pin_shared_region()
# All TP/PCP ranks hold identical encoder output. Only one rank
# per mmap needs to write — saves host memory bandwidth.
# DCP is a subdivision of TP, so tp_rank==0 covers it.
self._is_save_rank = (
get_tensor_model_parallel_rank() == 0 and get_pcp_group().rank_in_group == 0
)
# Descriptor buffer pool (recycled across steps, shared by both paths).
self._buf_pool = DescriptorBufferPool()
# Active save buffer being filled during save_caches calls this step.
self._save_bufs: DescriptorBuffers | None = None
self._save_count: int = 0
# Bytes accumulated into the active save buffer this step, for logging.
self._save_bytes: int = 0
# mm_hashes accumulated into the active save buffer this step.
self._save_mm_hashes: list[str] = []
# Stream the active save buffer's copy will run on, held from the moment
# the buffer opens so `save_caches` can register its sources against it.
self._save_stream: torch.Stream | None = None
# Batched copies whose end event has not yet fired. Buffers and stream
# stay held until the event fires; the completions are reported to the
# scheduler only then. Separate per direction because saves and loads
# can be in flight concurrently.
self._inflight_saves: deque[Transfer] = deque()
self._inflight_loads: deque[Transfer] = deque()
# Streams and timing events recycled across steps to avoid per-step
# create/destroy churn. Shared by both directions.
self._stream_pool: list[torch.Stream] = []
self._event_pool: list[torch.Event] = []
def _acquire_stream(self) -> torch.Stream:
if self._stream_pool:
return self._stream_pool.pop()
return current_platform.Stream()
def _acquire_event(self) -> torch.Event:
return (
self._event_pool.pop()
if self._event_pool
else torch.Event(enable_timing=True)
)
def _collect_finished(self, inflight: deque[Transfer], direction: str) -> list:
"""Pop transfers whose end event has fired, recycle their stream,
events, and buffers, and return their completions.
The front check is conservative: a transfer is reported only once its
own end event fires, so completions are always genuinely done, and each
is reported exactly once because the transfer is popped as it is
reported. A later transfer that happens to finish first simply waits
behind the front and is reported on a subsequent poll.
"""
done: list = []
while inflight and inflight[0].end_event.query():
transfer = inflight.popleft()
done.extend(transfer.completions)
elapsed_ms = transfer.start_event.elapsed_time(transfer.end_event)
logger.debug(
"EC %s: %d entr%s (%d bytes) took %.3f ms",
direction,
len(transfer.completions),
"y" if len(transfer.completions) == 1 else "ies",
transfer.num_bytes,
elapsed_ms,
)
self._buf_pool.release(transfer.bufs)
self._stream_pool.append(transfer.stream)
self._event_pool.append(transfer.start_event)
self._event_pool.append(transfer.end_event)
return done
def _pin_shared_region(self) -> None:
"""Prepare the shared region for device transfers."""
if is_pin_memory_available():
self._region.pin_memory()
def _submit_transfer(
self,
bufs: DescriptorBuffers,
count: int,
direction: ECCPUTransferDirection,
) -> None:
"""Submit a transfer on the current backend stream."""
src_ptrs = bufs.src_ptrs[:count]
dst_ptrs = bufs.dst_ptrs[:count]
sizes = bufs.sizes[:count]
if direction == ECCPUTransferDirection.HOST_TO_DEVICE:
swap_blocks_batch(
src_ptrs,
dst_ptrs,
sizes,
is_src_access_order_any=True,
)
else:
swap_blocks_batch(src_ptrs, dst_ptrs, sizes)
def _shutdown_transfer_backend(self) -> None:
"""Release backend-specific transfer resources."""
def save_caches(
self,
encoder_cache: dict[str, torch.Tensor],
mm_hash: str,
connector_metadata: ECCPUConnectorMetadata,
) -> None:
"""Fill descriptor buffers directly for batched flush.
Registers `encoder_cache[mm_hash]` against the stream the copy will run
on, so the caching allocator will not hand that memory to another
allocation before the copy has read it. The descriptor buffers hold raw
addresses, and the caller is free to evict the entry as soon as this
step's saves are dispatched.
"""
if not self._is_save_rank:
return
block_ids = connector_metadata.saves.get(mm_hash)
if block_ids is None:
return
src = encoder_cache[mm_hash]
total_bytes = src.numel() * src.element_size()
block_size = self._region.block_size_bytes
allocated_bytes = len(block_ids) * block_size
assert total_bytes <= allocated_bytes, (
f"EC: encoder output exceeds allocated blocks for mm_hash={mm_hash}: "
f"{total_bytes} bytes but only {allocated_bytes} allocated "
f"({len(block_ids)} blocks × {block_size}). "
f"shape={list(src.shape)}"
)
if self._save_bufs is None:
total = sum(len(v) for v in connector_metadata.saves.values())
self._save_bufs = self._buf_pool.acquire(total)
self._save_stream = self._acquire_stream()
self._save_mm_hashes = []
assert self._save_count + len(block_ids) <= self._save_bufs.src_ptrs.numel()
bufs = self._save_bufs
assert self._save_stream is not None
src.record_stream(self._save_stream)
# Device addresses are unsigned and can have the high bit set (XPU USM),
# so they are added up in uint64; signed arithmetic rejects them.
src_base = np.uint64(src.view(-1).view(torch.uint8).data_ptr())
dst_base = self._region.blocks.data_ptr()
# Copy only the tensor bytes, excluding any block padding.
n_used = -(-total_bytes // block_size)
if n_used:
slots, first_blocks, num_blocks = _coalesce_runs(block_ids[:n_used])
run_bytes = num_blocks * block_size
# The output can stop partway into the final run.
run_bytes[-1] = total_bytes - slots[-1] * block_size
bufs.add_copies(
self._save_count,
src_base + slots.astype(np.uint64) * np.uint64(block_size),
dst_base + first_blocks * block_size,
run_bytes,
)
self._save_count += slots.size
self._save_bytes += total_bytes
self._save_mm_hashes.append(mm_hash)
def flush_saves(self) -> None:
"""Flush all accumulated saves in a single swap_blocks_batch call.
Runs the copy on the stream the batch acquired when it opened, gated
behind the compute stream that produced the encoder outputs, brackets it
with start/end events, and enqueues the batch as in-flight. The stream,
events, and descriptor buffers are recycled once the end event fires (see
`_collect_finished`), which is also when the saved mm_hashes become
safe to mark ready.
"""
if self._save_count == 0:
return
bufs = self._save_bufs
stream = self._save_stream
assert bufs is not None and stream is not None
n = self._save_count
num_bytes = self._save_bytes
# Gate the GPU→CPU copy behind the compute stream: it reads the encoder
# outputs the model just produced there.
stream.wait_stream(current_platform.current_stream())
start_event = self._acquire_event()
end_event = self._acquire_event()
with current_platform.stream(stream):
start_event.record(stream)
self._submit_transfer(
bufs,
n,
ECCPUTransferDirection.DEVICE_TO_HOST,
)
end_event.record(stream)
self._inflight_saves.append(
Transfer(
start_event=start_event,
end_event=end_event,
completions=self._save_mm_hashes,
bufs=bufs,
stream=stream,
num_bytes=num_bytes,
)
)
self._save_bufs = None
self._save_stream = None
self._save_count = 0
self._save_bytes = 0
self._save_mm_hashes = []
def start_load_caches(
self,
encoder_cache: dict[str, torch.Tensor],
connector_metadata: ECCPUConnectorMetadata,
) -> None:
"""Consumer path: single batched copy of all loads from mmap→GPU.
The destination is registered against the compute stream that consumes
it, so the caching allocator will not hand that memory to a later load
while reads of it are still queued there.
"""
if not connector_metadata.loads:
return
block_size = self._region.block_size_bytes
blocks = self._region.blocks
dtype = self._dtype
device_type = current_platform.device_type
src_base = blocks.data_ptr()
# Copy every dispatched load. The scheduler routes an input here only
# after missing the GPU encoder cache, so a resident hit is not expected;
# copying unconditionally means every participating rank reports each
# transfer exactly once, which is what the scheduler counts to release
# the pin.
# Ordered by first block so entries that neighbour each other in the
# region land next to each other in the destination and coalesce into
# one descriptor. Every loop below consumes this same order.
load_items = sorted(
connector_metadata.loads.items(), key=lambda kv: kv[1][1][0]
)
total_blocks = sum(len(block_ids) for _, (_, block_ids, _) in load_items)
stream = self._acquire_stream()
compute_stream = current_platform.current_stream()
start_event = self._acquire_event()
end_event = self._acquire_event()
with current_platform.stream(stream):
# Single contiguous destination buffer for all loads.
dst_buf = torch.empty(
total_blocks, block_size, dtype=torch.int8, device=device_type
)
dst_buf.record_stream(compute_stream)
# See save_caches: device addresses are added up in uint64.
dst_buf_base = np.uint64(dst_buf.data_ptr())
bufs = self._buf_pool.acquire(total_blocks)
slots, first_blocks, num_blocks = _coalesce_runs(
list(chain.from_iterable(ids for _, (_, ids, _) in load_items))
)
bufs.add_copies(
0,
src_base + first_blocks * block_size,
dst_buf_base + slots.astype(np.uint64) * np.uint64(block_size),
num_blocks * block_size,
)
op_idx = slots.size
start_event.record(stream)
self._submit_transfer(
bufs,
op_idx,
ECCPUTransferDirection.HOST_TO_DEVICE,
)
end_event.record(stream)
self._inflight_loads.append(
Transfer(
start_event=start_event,
end_event=end_event,
completions=[tid for _, (tid, _, _) in load_items],
bufs=bufs,
stream=stream,
num_bytes=total_blocks * block_size,
)
)
# Slice contiguous buffer into per-hash views.
offset = 0
for mm_hash, (_, block_ids, shape) in load_items:
n = len(block_ids)
flat = dst_buf[offset : offset + n].view(dtype).reshape(-1)
encoder_cache[mm_hash] = flat[: shape[0] * shape[1]].view(shape)
offset += n
current_platform.current_stream().wait_stream(stream)
def build_connector_worker_meta(self) -> ECCPUWorkerMetadata | None:
"""Report the GPU copies that completed on this rank this step: saved
mm_hashes and loaded transfer ids.
Returns None when nothing finished, so the scheduler sees no payload.
"""
completed_saves = self._collect_finished(self._inflight_saves, "save")
completed_loads = self._collect_finished(self._inflight_loads, "load")
if not completed_saves and not completed_loads:
return None
return ECCPUWorkerMetadata(
completed_saves=completed_saves,
completed_loads=completed_loads,
)
def shutdown(self) -> None:
for transfer in (*self._inflight_saves, *self._inflight_loads):
transfer.end_event.synchronize()
self._shutdown_transfer_backend()
self._save_bufs = None
self._save_stream = None
self._save_count = 0
self._save_bytes = 0
self._save_mm_hashes = []
self._inflight_saves.clear()
self._inflight_loads.clear()
self._stream_pool.clear()
self._event_pool.clear()
try:
self._region.cleanup()
except Exception:
logger.debug("EC: worker region cleanup failed", exc_info=True)
vllm/distributed/ec_transfer/ec_connector/cpu/worker/descriptor_buffers.py
-90
View File
@@ -1,90 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Reusable pool of (src_ptrs, dst_ptrs, sizes) tensor triples.
Used by ECCPUWorker to batch swap_blocks_batch descriptors without
per-step allocation overhead.
"""
from typing import NamedTuple
import numpy as np
import torch
from vllm.platforms import current_platform
# CUDA/ROCm cache_kernels.cu requires int64 pointers; the XPU DMA engine
# requires uint64 (see vllm._custom_ops.swap_blocks_batch).
_PTR_DTYPE = torch.uint64 if current_platform.is_xpu() else torch.int64
class DescriptorBuffers(NamedTuple):
src_ptrs: torch.Tensor
dst_ptrs: torch.Tensor
sizes: torch.Tensor
# Numpy aliases of the three tensors, written via add_copies().
src_np: np.ndarray
dst_np: np.ndarray
sizes_np: np.ndarray
def add_copies(
self,
first: int,
sources: np.ndarray,
destinations: np.ndarray,
byte_counts: np.ndarray,
) -> None:
"""Store a batch of copies in these buffers, the first at *first*.
One copy is one source address, one destination address, and a number
of bytes. `swap_blocks_batch` performs every copy held in the buffers
in a single call, so a caller fills in all of a step's copies before
asking for any of them to run, and passes *first* to say where its own
copies begin.
Whole batches go in at once because storing a single number into a
torch tensor costs microseconds. Adding copies one at a time can take
longer than performing them.
TODO(torch>=2.14): write the tensors directly and drop these numpy
aliases once the minimum supported torch is 2.14. They exist only
because torch's setitem unpacks the value as a signed long long before
2.14 (pytorch#191458), rejecting XPU USM addresses >= 2**63, while the
two's-complement rewrite those versions accept is in turn rejected by
uint64 from 2.14 on. Numpy casts against the array's own type and so
works on either. Byte counts are small enough to need none of this
care, and use their alias only to match the addresses.
"""
last = first + len(sources)
self.src_np[first:last] = sources
self.dst_np[first:last] = destinations
self.sizes_np[first:last] = byte_counts
class DescriptorBufferPool:
"""Pool of descriptor buffer triples for swap_blocks_batch.
Each buffer is a `DescriptorBuffers` namedtuple of three 1-D tensors
(dtype `_PTR_DTYPE`, platform-dependent) of equal length, paired with
numpy aliases used to fill them. Buffers are recycled across steps; if a
returned buffer is too small it is discarded and a fresh one allocated.
"""
def __init__(self) -> None:
# LIFO stack of idle buffer triples.
self._pool: list[DescriptorBuffers] = []
def acquire(self, n: int) -> DescriptorBuffers:
"""Get a buffer triple with capacity >= *n*."""
if self._pool:
bufs = self._pool.pop()
if bufs.src_ptrs.numel() >= n:
return bufs
src, dst, sizes = (torch.empty(n, dtype=_PTR_DTYPE) for _ in range(3))
return DescriptorBuffers(
src, dst, sizes, src.numpy(), dst.numpy(), sizes.numpy()
)
def release(self, bufs: DescriptorBuffers) -> None:
"""Return a buffer triple to the pool for reuse."""
self._pool.append(bufs)
vllm/distributed/ec_transfer/ec_connector/example_connector.py
-231
View File
@@ -1,231 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import safetensors
from vllm.config import VllmConfig
from vllm.distributed.ec_transfer.ec_connector.base import (
ECConnectorBase,
ECConnectorMetadata,
ECConnectorRole,
)
from vllm.distributed.ec_transfer.ec_connector.utils import (
PlaceholderMetadataResolver,
collect_ec_item_metadata,
)
from vllm.logger import init_logger
from vllm.utils.gpu_sync_debug import gpu_sync_allowed
from vllm.v1.core.sched.output import SchedulerOutput
if TYPE_CHECKING:
from vllm.v1.request import Request
logger = init_logger(__name__)
@dataclass
class MMMeta:
mm_hash: str
num_token: int
@staticmethod
def make_meta(mm_hash, num_token) -> "MMMeta":
return MMMeta(mm_hash=mm_hash, num_token=num_token)
@dataclass
class ECExampleConnectorMetadata(ECConnectorMetadata):
mm_datas: list[MMMeta]
def __init__(self):
self.mm_datas = []
def add_mm_data(self, mm_data: MMMeta):
self.mm_datas.append(mm_data)
class ECExampleConnector(ECConnectorBase):
# NOTE: This is Simple debug implementation of the EC connector.
# It save / load the EC cache to / from the disk.
def __init__(self, vllm_config: "VllmConfig", role: ECConnectorRole):
super().__init__(vllm_config=vllm_config, role=role)
# req_id -> index
self._mm_datas_need_loads: dict[str, int] = {}
self._metadata_resolver = PlaceholderMetadataResolver(vllm_config.model_config)
transfer_config = vllm_config.ec_transfer_config
if transfer_config is not None:
self._storage_path = transfer_config.get_from_extra_config(
"shared_storage_path", "/tmp"
)
logger.debug(transfer_config)
logger.debug("Shared storage path is %s", self._storage_path)
else:
raise ValueError("ec_transfer_config must be set for ECConnectorBase")
def start_load_caches(self, encoder_cache, **kwargs) -> None:
"""Start loading the cache from the connector into vLLM's encoder cache.
This method loads the encoder cache based on metadata provided by the scheduler.
It is called before `_gather_mm_embeddings` for the EC Connector. For EC,
the `encoder_cache` and `mm_hash` are stored in `kwargs`.
Args:
encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
data hashes (`mm_hash`) to encoder cache tensors.
kwargs (dict): Additional keyword arguments for the connector.
"""
from vllm.platforms import current_platform
# Get the metadata
metadata: ECConnectorMetadata = self._get_connector_metadata()
assert isinstance(metadata, ECExampleConnectorMetadata)
assert encoder_cache is not None
if metadata is None:
logger.warning(
"In connector.start_load_caches, but the connector metadata is None"
)
return
# A bare "cuda" makes safetensors load onto cuda:0. Pin the load to the
# selected device so tensor-parallel workers load the cache locally.
device = current_platform.device_type
if current_platform.is_cuda_alike():
device = f"{device}:{current_platform.current_device()}"
# Load the EC for each mm data
for mm_data in metadata.mm_datas:
if mm_data.mm_hash in encoder_cache:
continue
filename = self._generate_filename_debug(mm_data.mm_hash)
with gpu_sync_allowed():
ec_cache = safetensors.torch.load_file(filename, device=device)[
"ec_cache"
]
encoder_cache[mm_data.mm_hash] = ec_cache
logger.debug("Success load encoder cache for hash %s", mm_data.mm_hash)
def save_caches(self, encoder_cache, mm_hash, **kwargs) -> None:
"""Save the encoder cache to the connector.
This method saves the encoder cache from the worker's local storage
to shared storage or another external connector.
Args:
encoder_cache (dict[str, torch.Tensor]): A dictionary mapping multimodal
data hashes (`mm_hash`) to encoder cache tensors.
mm_hash (str): The hash of the multimodal data whose cache is being saved.
kwargs (dict): Additional keyword arguments for the connector.
"""
# Return if it is PD Instance
if not self.is_producer:
return
filename = self._generate_filename_debug(mm_hash)
ec_cache = encoder_cache[mm_hash]
with gpu_sync_allowed():
tensors = {"ec_cache": ec_cache.detach().cpu()}
safetensors.torch.save_file(tensors, filename)
logger.debug("Save cache successful for mm_hash %s", mm_hash)
def has_cache_item(
self,
identifier: str,
) -> bool:
"""Check if cache exist externally for the media.
Args:
identifier (str): the identifier of the media.
Returns:
Bool indicate that media exists in cache or not
"""
return self._found_match_for_mm_data(identifier)
def update_state_after_alloc(
self,
request: "Request",
index: int,
) -> None:
"""Update ECConnector state after encoder cache allocation."""
mm_hash = request.mm_features[index].identifier
# Only load cache if it is consumer and cache exists
if not self.is_consumer or not self.has_cache_item(mm_hash):
return
num_encoder_token = request.get_num_encoder_embeds(index)
self._mm_datas_need_loads[mm_hash] = num_encoder_token
def build_connector_meta(
self,
scheduler_output: SchedulerOutput,
) -> ECConnectorMetadata:
"""Build the connector metadata for this step.
This function should NOT modify any fields in the scheduler_output.
Also, calling this function will reset the state of the connector.
This only build for load mm_data only
Args:
scheduler_output (SchedulerOutput): the scheduler output object.
"""
meta = ECExampleConnectorMetadata()
for mm_hash, num_encoder_token in self._mm_datas_need_loads.items():
meta.add_mm_data(MMMeta.make_meta(mm_hash, num_encoder_token))
self._mm_datas_need_loads.clear()
return meta
def request_finished(
self,
request: "Request",
) -> tuple[bool, dict[str, Any] | None]:
"""Report each item's cache key and grid so a consumer can skip the
image transform.
A consumer only needs the grid to size the prompt's placeholder range;
the embedding itself arrives through this connector. Reporting the grid
the producer actually computed keeps the two sides in agreement without
the caller re-deriving it from the raw media.
"""
if not self.is_producer:
return False, None
items = collect_ec_item_metadata(request.mm_features, self._metadata_resolver)
if not items:
return False, None
return False, items
# ==============================
# Helper functions
# ==============================
def _found_match_for_mm_data(self, mm_hash) -> bool:
"""Check if the cache is hit for the request."""
filename = self._generate_filename_debug(mm_hash, create_folder=False)
return os.path.exists(filename)
def _generate_foldername_debug(
self,
mm_hash: str,
create_folder: bool = True, # <- now defaults to True
) -> str:
"""Return the folder in which the cache for this mm_hash lives.
If `create_folder` is True (default) the directory is created
recursively the first time it is needed.
"""
foldername = os.path.join(self._storage_path, mm_hash)
if create_folder:
os.makedirs(foldername, exist_ok=True)
return foldername
def _generate_filename_debug(
self,
mm_hash: str,
create_folder: bool = True,
) -> str:
"""Return the full path of the safetensors file for this mm_hash.
If `create_folder` is True (default) the parent directory is created.
"""
foldername = self._generate_foldername_debug(mm_hash, create_folder)
return os.path.join(foldername, "encoder_cache.safetensors")
vllm/entrypoints/cli/main.py
+2 -10
View File
@@ -47,7 +47,6 @@ def main():
import vllm.entrypoints.cli.preload
import vllm.entrypoints.cli.run_batch
import vllm.entrypoints.cli.serve
import vllm.entrypoints.cli.snapshot
from vllm.entrypoints.serve.utils.api_utils import (
VLLM_SUBCMD_PARSER_EPILOG,
cli_env_setup,
@@ -62,11 +61,9 @@ def main():
vllm.entrypoints.cli.download_kernels,
vllm.entrypoints.cli.preload,
vllm.entrypoints.cli.run_batch,
vllm.entrypoints.cli.snapshot,
]
if sys.argv[1:2] != ["snapshot"]:
cli_env_setup()
cli_env_setup()
parser = FlexibleArgumentParser(
description="vLLM CLI",
@@ -81,12 +78,7 @@ def main():
subparsers = parser.add_subparsers(required=False, dest="subparser")
cmds = {}
for cmd_module in CMD_MODULES:
if cmd_module is vllm.entrypoints.cli.snapshot:
new_cmds = cmd_module.cmd_init(
create_requested=sys.argv[1:3] == ["snapshot", "create"]
)
else:
new_cmds = cmd_module.cmd_init()
new_cmds = cmd_module.cmd_init()
for cmd in new_cmds:
cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)
cmds[cmd.name] = cmd
vllm/entrypoints/cli/snapshot.py
-143
View File
@@ -1,143 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
import os
import platform
import typing
from pathlib import Path
from vllm.entrypoints.cli.types import CLISubcommand
if typing.TYPE_CHECKING:
from vllm.utils.argparse_utils import FlexibleArgumentParser
else:
FlexibleArgumentParser = argparse.ArgumentParser
def validate_create_args(args: argparse.Namespace) -> None:
model = str(getattr(args, "model_tag", None) or args.model)
revision = str(getattr(args, "revision", None) or "")
if (
Path(model).exists()
or len(revision) != 40
or any(character not in "0123456789abcdefABCDEF" for character in revision)
):
raise ValueError(
"snapshot create requires an immutable remote model and "
"40-character --revision"
)
parallel_sizes = (
("tensor parallel", args.tensor_parallel_size),
("pipeline parallel", args.pipeline_parallel_size),
("data parallel", args.data_parallel_size),
("prefill context parallel", args.prefill_context_parallel_size),
)
for name, size in parallel_sizes:
if size != 1:
raise ValueError(f"snapshot create currently requires {name} size 1")
unsupported_frontend = (
args.api_server_count not in (None, 1)
or args.grpc
or args.headless
or args.api_key
or os.environ.get("VLLM_API_KEY")
or args.uds
or args.middleware
)
if unsupported_frontend or any(
(
args.ssl_keyfile,
args.ssl_certfile,
args.ssl_ca_certs,
args.enable_ssl_refresh,
args.ssl_ciphers,
)
):
raise ValueError(
"snapshot create requires one unauthenticated plaintext HTTP server"
)
if args.logprobs_mode in {"raw_logits", "processed_logits"}:
raise ValueError("snapshot canary requires a log-probability mode")
if any(
getattr(args, name, None) is not None
for name in ("speculative_config", "spec_method", "spec_model", "spec_tokens")
):
raise ValueError("snapshot create does not support speculative decoding")
if platform.system() != "Linux" or platform.machine() != "x86_64":
raise ValueError("snapshot create requires Linux x86_64")
def run_create(args: argparse.Namespace) -> None:
validate_create_args(args)
from vllm.snapshot.controller import create_snapshot
create_snapshot(args)
def run_inspect(args: argparse.Namespace) -> None:
from vllm.snapshot.manifest import inspect_snapshot
print(json.dumps(inspect_snapshot(Path(args.snapshot_dir)), indent=2))
def run_restore(args: argparse.Namespace) -> None:
from vllm.snapshot.controller import restore_snapshot
restore_snapshot(args)
class SnapshotSubcommand(CLISubcommand):
"""The `snapshot` subcommand for the vLLM CLI."""
name = "snapshot"
def __init__(self, *, create_requested: bool = False) -> None:
self.create_requested = create_requested
@staticmethod
def cmd(args: argparse.Namespace) -> None:
args.snapshot_dispatch(args)
def subparser_init(
self, subparsers: argparse._SubParsersAction
) -> FlexibleArgumentParser:
parser = subparsers.add_parser(
self.name,
help="Create, inspect, or restore an initialized vLLM snapshot.",
usage="vllm snapshot <create|inspect|restore> [options]",
)
actions = parser.add_subparsers(required=True, dest="snapshot_action")
create_parser = actions.add_parser(
"create", help="Create a snapshot from an initialized TP1 engine."
)
if self.create_requested:
from vllm.entrypoints.launchers.cli_args import make_arg_parser
create_parser = make_arg_parser(create_parser)
else:
create_parser.add_argument("model_tag", nargs="?")
create_parser.add_argument("--snapshot-dir", required=True)
create_parser.set_defaults(snapshot_dispatch=run_create)
inspect_parser = actions.add_parser(
"inspect", help="Inspect a snapshot without restoring it."
)
inspect_parser.add_argument("snapshot_dir")
inspect_parser.set_defaults(snapshot_dispatch=run_inspect)
restore_parser = actions.add_parser(
"restore", help="Restore a same-host TP1 snapshot."
)
restore_parser.add_argument("snapshot_dir")
restore_parser.add_argument("--host", default=None)
restore_parser.add_argument("--port", type=int, default=8000)
restore_parser.set_defaults(snapshot_dispatch=run_restore)
return parser
def cmd_init(*, create_requested: bool = False) -> list[CLISubcommand]:
return [SnapshotSubcommand(create_requested=create_requested)]
vllm/model_executor/hw_agnostic/__init__.py
-2
View File
@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/hw_agnostic/custom_op.py
-314
View File
@@ -1,314 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import functools
import inspect
import torch
import torch.nn as nn
from vllm.config import get_cached_compilation_config
from vllm.logger import init_logger
from vllm.model_executor.utils import maybe_disable_graph_partition
from vllm.platforms import current_platform
logger = init_logger(__name__)
# Dictionary of all custom ops (classes, indexed by registered name).
# To check if an op with a name is enabled, call .enabled() on the class.
# Examples:
# - MyOp.enabled()
# - op_registry["my_op"].enabled()
op_registry: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
op_registry_oot: dict[str, type["CustomOp"] | type["PluggableLayer"]] = {}
def maybe_get_oot_by_class(class_type: type) -> type:
class_name = class_type.__name__
if class_name in op_registry_oot:
return op_registry_oot[class_name]
return class_type
class PluggableLayer(nn.Module):
"""Base class for pluggable layers.
A PluggableLayer is a *module-composing* abstraction: it may instantiate other
``torch.nn.Module`` objects as sub-layers, and its functionality depends on
these sub-layers following a generalized invocation sequence. Also, it is stateful
and may hold parameters or buffers.
Unlike :class:`CustomOp`, PluggableLayer does NOT provide per-platform
``forward_*`` dispatch. Instead, it supports out-of-tree (OOT) replacement
of the entire layer class at instantiation time, allowing customized
initialization and submodule composition.
"""
def __new__(cls, *args, **kwargs):
try:
layer_class_name = cls.__name__
except AttributeError:
raise TypeError(
f"Cannot instantiate '{cls.__name__}': its 'name' attribute "
f"was not set, possibly because it was not decorated with "
f"@PluggableLayer.register, or it's the PluggableLayer itself."
) from None
if layer_class_name not in op_registry_oot:
layer_cls_to_instantiate = cls
else:
layer_cls_to_instantiate = op_registry_oot[layer_class_name]
logger.debug(
"Instantiating pluggable layer: %s using %s",
layer_class_name,
str(layer_cls_to_instantiate),
)
return super().__new__(layer_cls_to_instantiate)
# Decorator to register pluggable layers.
@classmethod
def register(cls, name: str):
def decorator(op_cls):
assert name not in op_registry, f"Duplicate op name: {name}"
op_cls.name = name
op_registry[name] = op_cls
return op_cls
return decorator
# Decorator to register out-of-tree(oot) pluggable layers.
# For OOT pluggable layers:
# if in-tree layer class is registered with an oot_custom_layer,
# the oot_custom_layer will be used instead.
@classmethod
def register_oot(cls, _decorated_layer_cls=None, name: str | None = None):
def decorator(layer_cls):
reg_name = name if name is not None else cls.__name__
assert reg_name not in op_registry_oot, f"Duplicate layer name: {reg_name}"
layer_cls.name = reg_name
op_registry_oot[reg_name] = layer_cls
return layer_cls
if _decorated_layer_cls is None:
# Called with parentheses: @PluggableLayer.register_oot()
# or @PluggableLayer.register_oot(name="...")
return decorator
elif isinstance(_decorated_layer_cls, type): # Check if it's a class
# Called without parentheses: @PluggableLayer.register_oot
return decorator(_decorated_layer_cls)
else:
raise TypeError("Decorator can only be applied to classes.")
class CustomOp(nn.Module):
"""Base class for custom ops.
Dispatches the forward method to the appropriate backend.
"""
def __new__(cls, *args, **kwargs):
try:
op_name = cls.__name__
except AttributeError:
raise TypeError(
f"Cannot instantiate '{cls.__name__}': its 'name' attribute "
f"was not set, possibly because it was not decorated with "
f"@CustomOp.register, or it's the CustomOp base class itself."
) from None
if op_name not in op_registry_oot:
op_cls_to_instantiate = cls
else:
op_cls_to_instantiate = op_registry_oot[op_name]
logger.debug(
"Instantiating custom op: %s using %s",
op_name,
str(op_cls_to_instantiate),
)
return super().__new__(op_cls_to_instantiate)
def __init__(self, *, enforce_enable: bool = False, compile_native: bool = False):
super().__init__()
self._enforce_enable = enforce_enable
self._forward_method = self.dispatch_forward(compile_native=compile_native)
def forward(self, *args, **kwargs):
return self._forward_method(*args, **kwargs)
def forward_native(self, *args, **kwargs):
"""PyTorch-native implementation; OOT plugins override via
:meth:`forward_oot`."""
raise NotImplementedError
def forward_oot(self, *args, **kwargs):
# By default, OOT ops fall back to the PyTorch-native implementation.
return self.forward_native(*args, **kwargs)
def dispatch_forward(self, compile_native: bool):
# The hw-agnostic CustomOp dispatches only between forward_native
# (in-tree, used on every platform) and forward_oot (overridden by
# an OOT plugin via register_oot).
compilation_config = get_cached_compilation_config()
# NOTE(shen-shanshan): CustomOp object can be enforce enabled, e.g.,
# enable device-specific kernels in ViT models when enabling graph
# mode. By default, it will follow the compilation_config to determine
# whether enable itself.
# This enforce_enable mechanism will be removed after we adding a
# separate compilation_config for multi-modal part.
enabled = self._enforce_enable or self.enabled()
if enabled:
compilation_config.enabled_custom_ops.update([self.__class__.name])
else:
compilation_config.disabled_custom_ops.update([self.__class__.name])
if not enabled:
# Compile forward_native to avoid eager torch ops if inside
# opaque torch custom op (e.g. fused_moe, unified_attention, etc.)
return self.maybe_compile(self.forward_native, enable=compile_native)
if current_platform.is_out_of_tree():
return self.forward_oot
return self.forward_native
def maybe_compile(self, fn, *, enable: bool = True):
"""Compile fn if compilation enabled.
Useful for CustomOp instances called from within a torch custom op,
meaning the forward call is hidden from the model-level torch.compile.
NOTE: this does not enable fusion across ops, so opaque custom ops
should still be unwrapped wherever possible.
"""
from vllm.config.compilation import CompilationMode
# Do not compile if compilation disabled
if not enable:
return fn
# Do not compile if global compilation disabled
compilation_config = get_cached_compilation_config()
if compilation_config.mode == CompilationMode.NONE:
return fn
# If eager backend is used, do not compile either
if compilation_config.backend == "eager":
return fn
compile_options = maybe_disable_graph_partition(
current_platform.simple_compile_backend
)
backend = current_platform.simple_compile_backend
dynamic_arg_dims = getattr(self.__class__, "_dynamic_arg_dims", None)
if dynamic_arg_dims is not None:
compiled_fn = torch.compile(
fn,
dynamic=False,
backend=backend,
options=compile_options,
)
sig = inspect.signature(fn)
@functools.wraps(fn)
def wrapper(*args, **kwargs):
bound = sig.bind(*args, **kwargs)
bound.apply_defaults()
for name, dims in dynamic_arg_dims.items():
arg = bound.arguments.get(name)
if arg is not None and isinstance(arg, torch.Tensor):
dims_list = [dims] if isinstance(dims, int) else dims
for d in dims_list:
real_d = arg.ndim + d if d < 0 else d
torch._dynamo.mark_dynamic(arg, real_d)
return compiled_fn(*args, **kwargs)
return wrapper
# dynamic=True to avoid recompilations
return torch.compile(
fn,
dynamic=True,
backend=backend,
options=compile_options,
)
@classmethod
def enabled(cls) -> bool:
# if no name, then it was not registered
compilation_config = get_cached_compilation_config()
custom_ops = compilation_config.custom_ops
if not hasattr(cls, "name"):
logger.warning_once(
"Custom op %s was not registered, which means it won't appear "
"in the op registry. It will be enabled/disabled based on the "
"global settings.",
cls.__name__,
)
return CustomOp.default_on()
enabled = f"+{cls.name}" in custom_ops
disabled = f"-{cls.name}" in custom_ops
assert not (enabled and disabled), f"Cannot enable and disable {cls.name}"
return (CustomOp.default_on() or enabled) and not disabled
@staticmethod
def default_on() -> bool:
"""Behavior controlled by `CompilationConfig.custom_ops`: On by default if
'all', off by default if 'none'.
When PyTorch Inductor is used, 'none' is the default value,
otherwise 'all'.
"""
compilation_config = get_cached_compilation_config()
count_none = compilation_config.custom_ops.count("none")
count_all = compilation_config.custom_ops.count("all")
assert count_none + count_all == 1
return not count_none > 0 or count_all > 0
# Decorator to register custom ops.
@classmethod
def register(
cls,
name: str,
dynamic_arg_dims: dict[str, int | list[int]] | None = None,
):
def decorator(op_cls):
assert name not in op_registry, f"Duplicate op name: {name}"
op_cls.name = name
op_cls._dynamic_arg_dims = dynamic_arg_dims
op_registry[name] = op_cls
return op_cls
return decorator
# Decorator to register out-of-tree(oot) custom ops.
# For OOT custom ops:
# if in-tree layer class is registered with an oot_custom_op layer,
# the oot_custom_op layer will be used instead.
# Example:
# - @UnquantizedFusedMoEMethod.register_oot
# class HPUUnquantizedFusedMoEMethod(UnquantizedFusedMoEMethod)
# or
# - @CustomOP.register_oot(name="UnquantizedFusedMoEMethod")
@classmethod
def register_oot(cls, _decorated_op_cls=None, name: str | None = None):
def decorator(op_cls):
reg_name = name if name is not None else cls.__name__
assert reg_name not in op_registry_oot, f"Duplicate op name: {reg_name}"
op_cls.name = reg_name
op_registry_oot[reg_name] = op_cls
return op_cls
if _decorated_op_cls is None:
# Called with parentheses: @CustomOP.register_oot()
# or @CustomOP.register_oot(name="...")
# So, _decorated_op_cls is None.
# We return the actual decorator function.
return decorator
elif isinstance(_decorated_op_cls, type): # Check if it's a class
# Called without parentheses: @CustomOP.register_oot
# The first argument is the class itself.
# We call the 'decorator' function immediately with the class.
return decorator(_decorated_op_cls)
else:
# Handle other unexpected cases if necessary
raise TypeError("Decorator can only be applied to classes.")
vllm/model_executor/hw_agnostic/layers/__init__.py
-2
View File
@@ -1,2 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
vllm/model_executor/hw_agnostic/layers/activation.py
-27
View File
@@ -1,27 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import torch.nn.functional as F
from vllm.model_executor.hw_agnostic.custom_op import CustomOp
@CustomOp.register("silu_and_mul")
class SiluAndMul(CustomOp):
"""SwiGLU: ``x -> silu(x[:d]) * x[d:]`` where ``d = x.shape[-1] // 2``."""
def forward_native(self, x: torch.Tensor) -> torch.Tensor:
d = x.shape[-1] // 2
return F.silu(x[..., :d]) * x[..., d:]
# Activation-and-mul ops keyed by HF activation name.
_ACTIVATION_AND_MUL_REGISTRY = {
"silu": SiluAndMul,
"swish": SiluAndMul,
}
def get_act_and_mul_fn(act_fn_name: str) -> CustomOp:
"""Build the hw-agnostic activation-and-mul op named `act_fn_name`."""
return _ACTIVATION_AND_MUL_REGISTRY[act_fn_name.lower()]()
vllm/model_executor/hw_agnostic/layers/layernorm.py
-90
View File
@@ -1,90 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import torch
import torch.nn as nn
from torch import Tensor
from vllm.model_executor.hw_agnostic.custom_op import CustomOp
@CustomOp.register("rms_norm")
class RMSNorm(CustomOp):
"""``x -> w * x / sqrt(E[x^2] + eps)``. With ``residual``, fuses
``residual += x`` then RMSNorm and returns ``(normalized, residual)``."""
def __init__(
self,
hidden_size: int,
eps: float = 1e-6,
var_hidden_size: int | None = None,
has_weight: bool = True,
dtype: torch.dtype | None = None,
) -> None:
super().__init__()
self.hidden_size = hidden_size
self.variance_epsilon = eps
self.variance_size_override = (
None if var_hidden_size == hidden_size else var_hidden_size
)
weight_dtype = dtype or torch.get_default_dtype()
self.has_weight = has_weight
weight = torch.ones(hidden_size, dtype=weight_dtype)
if has_weight:
self.weight = nn.Parameter(weight)
else:
self.register_buffer("weight", weight, persistent=False)
def _rms_norm(
self,
x: Tensor,
weight: Tensor | None,
epsilon: float,
variance_size: int | None = None,
) -> Tensor:
"""Weighted root-mean-square layer normalization"""
orig_dtype = x.dtype
x = x.to(torch.float32)
x_var = x if variance_size is None else x[..., :variance_size]
variance = x_var.pow(2).mean(dim=-1, keepdim=True)
x = x * torch.rsqrt(variance + epsilon)
if weight is not None:
x = x.to(weight.dtype) * weight
return x.to(orig_dtype)
def _fused_add_rms_norm(
self,
x: Tensor,
x_residual: Tensor,
weight: Tensor | None,
epsilon: float,
variance_size: int | None = None,
) -> tuple[Tensor, Tensor]:
"""Fused add and weighted root-mean-square layer normalization"""
orig_dtype = x.dtype
x = x.to(torch.float32)
x = x + x_residual.to(torch.float32)
x_residual = x.to(orig_dtype)
x_var = x if variance_size is None else x[..., :variance_size]
variance = x_var.pow(2).mean(dim=-1, keepdim=True)
x = x * torch.rsqrt(variance + epsilon)
if weight is not None:
x = x.to(weight.dtype) * weight
return x.to(orig_dtype), x_residual
def forward_native(
self,
x: torch.Tensor,
residual: torch.Tensor | None = None,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
weight = self.weight if self.has_weight else None
epsilon = self.variance_epsilon
variance_size = self.variance_size_override
if residual is None:
return self._rms_norm(x, weight, epsilon, variance_size)
else:
return self._fused_add_rms_norm(x, residual, weight, epsilon, variance_size)
def extra_repr(self) -> str:
return f"hidden_size={self.hidden_size}, eps={self.variance_epsilon}"
vllm/model_executor/layers/fused_allreduce_gemma_rms_norm.py
-190
View File
@@ -1,190 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Manual fusion of tensor-parallel all-reduce with the following GemmaRMSNorm.
Under tensor parallelism a ``RowParallelLinear`` (e.g. attention ``o_proj``)
produces a per-rank partial sum that is all-reduced, and the result is then fed
into a ``GemmaRMSNorm`` that adds the residual and normalizes. flashinfer ships a
kernel that fuses all-reduce + residual-add + RMSNorm into a single launch; this
helper drives it directly (no torch.compile pass) for models that run eager.
Scope: attention output only, no quantization. When the flashinfer fast path is
not applicable (TP==1, flashinfer/NVSwitch unavailable, unsupported dtype, or an
oversize batch) it falls back to ``all_reduce`` + ``GemmaRMSNorm``, which is
numerically identical to the unfused model path.
"""
import torch
from vllm.distributed.communication_op import tensor_model_parallel_all_reduce
from vllm.distributed.parallel_state import (
get_tensor_model_parallel_rank,
get_tensor_model_parallel_world_size,
get_tp_group,
)
from vllm.model_executor.layers.layernorm import GemmaRMSNorm
from vllm.platforms import current_platform
MiB = 1024 * 1024
# flashinfer fused all-reduce + RMSNorm is wired as a registered custom op in
# allreduce_rms_fusion; both that op and the workspace helpers only exist when
# flashinfer.comm.allreduce_fusion is importable.
try:
from vllm.compilation.passes.fusion.allreduce_rms_fusion import (
flashinfer_trtllm_fused_allreduce_norm,
)
from vllm.distributed.device_communicators.flashinfer_all_reduce import (
flashinfer_comm,
get_fi_ar_workspace,
)
_AR_RESIDUAL_RMS_NORM = (
flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm
if flashinfer_comm is not None
else None
)
except (ImportError, AttributeError):
flashinfer_trtllm_fused_allreduce_norm = None # type: ignore[assignment]
get_fi_ar_workspace = None # type: ignore[assignment]
_AR_RESIDUAL_RMS_NORM = None
_FI_SUPPORTED_DTYPES = (torch.bfloat16, torch.float16)
def _max_token_num(tp_size: int, hidden_size: int, dtype: torch.dtype) -> int | None:
"""Workspace token budget for flashinfer fused all-reduce, or None if the
current world size / device is unsupported. Mirrors ``FlashInferAllReduce``."""
from vllm.config.compilation import PassConfig
max_size_mb = PassConfig.default_fi_allreduce_fusion_max_size_mb().get(tp_size)
if not max_size_mb:
return None
element_size = torch.tensor([], dtype=dtype).element_size()
return int(max_size_mb * MiB) // (hidden_size * element_size)
def _can_use_flashinfer(hidden_states: torch.Tensor, tp_size: int) -> tuple[bool, int]:
"""Whether the flashinfer fused path applies; returns (ok, max_token_num)."""
if (
flashinfer_trtllm_fused_allreduce_norm is None
or get_fi_ar_workspace is None
or _AR_RESIDUAL_RMS_NORM is None
):
return False, 0
if (
not hidden_states.is_cuda
or hidden_states.dim() != 2
or not hidden_states.is_contiguous()
or hidden_states.dtype not in _FI_SUPPORTED_DTYPES
):
return False, 0
num_tokens, hidden_size = hidden_states.shape
max_token_num = _max_token_num(tp_size, hidden_size, hidden_states.dtype)
if max_token_num is None or num_tokens > max_token_num:
return False, 0
# Lazily create / fetch the (globally cached) workspace; returns None on
# GPUs without NVSwitch, in which case we fall back gracefully.
workspace = get_fi_ar_workspace(
world_size=tp_size,
rank=get_tensor_model_parallel_rank(),
max_token_num=max_token_num,
hidden_dim=hidden_size,
dtype=hidden_states.dtype,
group=get_tp_group().cpu_group,
)
if workspace is None:
return False, 0
# The token-count bound above uses the whole workspace budget, but a backend
# may use only a fraction of it per call (mnnvl rotates through three Lamport
# buffers). Ask the workspace so we don't admit tensors that the kernel will
# reject with "The buffer size in the given workspace is insufficient".
if not workspace.is_buffer_size_sufficient(
tp_size=tp_size,
num_tokens=num_tokens,
hidden_dim=hidden_size,
dtype=hidden_states.dtype,
):
return False, 0
return True, max_token_num
def _can_use_aiter_fused_ar_rms(hidden_states: torch.Tensor) -> bool:
if not current_platform.is_rocm():
return False
from vllm._aiter_ops import rocm_aiter_ops
if not rocm_aiter_ops.is_custom_all_reduce_enabled():
return False
if (
hidden_states.dim() != 2
or not hidden_states.is_contiguous()
or hidden_states.dtype not in _FI_SUPPORTED_DTYPES
):
return False
aiter_ar = rocm_aiter_ops.get_aiter_allreduce()
if aiter_ar is None or aiter_ar.disabled:
return False
total_bytes = hidden_states.numel() * hidden_states.element_size()
if total_bytes > aiter_ar.effective_max_size():
return False
if not aiter_ar.should_custom_ar(hidden_states):
return False
return aiter_ar.use_1stage_fused_ar_rms(hidden_states)
def fused_allreduce_gemma_rms_norm(
hidden_states: torch.Tensor,
residual: torch.Tensor,
norm: GemmaRMSNorm,
) -> tuple[torch.Tensor, torch.Tensor]:
"""All-reduce ``hidden_states`` + add ``residual`` + GemmaRMSNorm, fused.
``hidden_states`` is the per-rank *partial* (un-reduced) output of a
row-parallel linear; ``norm`` is the GemmaRMSNorm applied right after.
Returns ``(normed_output, new_residual)``, equivalent to
``norm(all_reduce(hidden_states), residual)``.
"""
tp_size = get_tensor_model_parallel_world_size()
if tp_size == 1:
# No all-reduce needed; identical to the unfused path.
return norm(hidden_states, residual)
ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size)
if ok:
norm_out = torch.empty_like(hidden_states)
# With norm_out provided, the kernel writes the new residual
# (all_reduce(hidden_states) + residual) into the hidden_states buffer
# and the normalized result into norm_out, leaving `residual` untouched.
flashinfer_trtllm_fused_allreduce_norm(
allreduce_in=hidden_states,
residual=residual,
rms_gamma=norm.weight,
rms_eps=norm.variance_epsilon,
world_size=tp_size,
weight_bias=1.0, # GemmaRMSNorm-style
launch_with_pdl=True,
fp32_acc=True,
max_token_num=max_token_num,
pattern_code=_AR_RESIDUAL_RMS_NORM,
norm_out=norm_out,
)
return norm_out, hidden_states
if _can_use_aiter_fused_ar_rms(hidden_states):
from vllm._aiter_ops import rocm_aiter_ops
return rocm_aiter_ops.get_fused_allreduce_rmsnorm_op()(
input_=hidden_states,
residual=residual,
weight=norm.weight,
epsilon=norm.variance_epsilon,
gemma_norm=True,
)
# Fallback: explicit all-reduce + GemmaRMSNorm (matches the unfused model).
reduced = tensor_model_parallel_all_reduce(hidden_states)
return norm(reduced, residual)
vllm/model_executor/layers/fused_moe/deep_gemm_mega_moe.py
-857
View File
@@ -1,857 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""DeepGEMM MegaMoE experts: fused EP dispatch, expert GEMMs and combine.
Each ``MEGA_MOE_*`` format is one row of the SM100 feature matrix. MX weights
and activations carry one UE8M0 scale per 1x32 block. NVFP4 weights carry one
E4M3 scale per 1x16 block plus a per-expert FP32 global scale; NVFP4 dispatch
scales each token by its own FP32 scale.
"""
import math
from collections.abc import Callable
from dataclasses import dataclass
from inspect import signature
from typing import Any, ClassVar, Literal
import torch
import torch.nn as nn
import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.distributed import get_ep_group
from vllm.distributed.eplb.eplb_state import EplbLayerState
from vllm.forward_context import get_forward_context, is_forward_context_available
from vllm.logger import init_logger
from vllm.model_executor.layers.fused_moe.prepare_megamoe import (
prepare_megamoe_inputs,
prepare_megamoe_nvfp4_inputs,
)
from vllm.model_executor.layers.fused_moe.router.base_router import (
eplb_map_to_physical_and_record,
)
from vllm.model_executor.layers.quantization.utils.quant_utils import (
QuantKey,
kMxfp8Dynamic,
kNvfp4DynamicToken,
)
from vllm.model_executor.models.utils import extract_layer_index
from vllm.model_executor.utils import set_weight_attrs
from vllm.platforms import current_platform
from vllm.v1.worker.ubatching import dbo_current_ubatch_id
logger = init_logger(__name__)
@dataclass(frozen=True)
class MegaMoEActivation:
"""Gated activation fused into the MegaMoE L1 epilogue.
``swiglu`` computes ``gate * sigmoid(alpha * gate) * (up + beta)`` after
clamping gate from above and up symmetrically to ``clamp``; this covers
plain, DeepSeek-V4 (``clamp``) and GPT-OSS style (``alpha``/``beta``)
experts. ``situ`` is Kimi K3's SiTU with gate tanh scale ``alpha`` and
linear tanh scale ``beta``; ``beta=0`` leaves the linear branch unscaled.
"""
name: Literal["swiglu", "situ"] = "swiglu"
clamp: float | None = None
alpha: float = 1.0
beta: float = 0.0
def __post_init__(self) -> None:
if self.name not in ("swiglu", "situ"):
raise ValueError(f"Unsupported MegaMoE activation {self.name!r}.")
if not (math.isfinite(self.alpha) and math.isfinite(self.beta)):
raise ValueError(f"MegaMoE activation scales must be finite: {self}.")
if self.clamp is not None and not self.clamp >= 0:
raise ValueError(f"MegaMoE activation clamp must be >= 0: {self}.")
if self.name == "situ" and (
self.clamp is not None or self.alpha <= 0 or self.beta < 0
):
raise ValueError(
"SiTU takes no clamp, a positive alpha and a non-negative beta; "
f"got {self}."
)
def kernel_kwargs(self) -> dict[str, Any]:
# Defaults stay implicit so DeepGEMM builds without alpha/beta still work.
kwargs: dict[str, Any] = {"activation_clamp": self.clamp}
if self.name != "swiglu":
kwargs["activation"] = self.name
if self.alpha != 1.0:
kwargs["activation_alpha"] = self.alpha
if self.beta != 0.0:
kwargs["activation_beta"] = self.beta
return kwargs
@dataclass(frozen=True)
class MegaMoEWeightSpec:
"""Storage of one MegaMoE weight in the kernel layout."""
dtype: torch.dtype
# Weight values packed per storage element along K (2 for FP4).
values_per_element: int = 1
# K values per block scale; None for unscaled weights.
scale_group_k: int | None = None
# Block scale dtype: UE8M0 (MX) or E4M3 (NVFP4).
scale_dtype: torch.dtype | None = None
@property
def has_global_scale(self) -> bool:
"""E4M3 block scales pair with an FP32 per-tensor scale."""
return self.scale_dtype == torch.float8_e4m3fn
_MXFP4_WEIGHT = MegaMoEWeightSpec(
torch.int8,
values_per_element=2,
scale_group_k=32,
scale_dtype=torch.float8_e8m0fnu,
)
_MXFP8_WEIGHT = MegaMoEWeightSpec(
torch.float8_e4m3fn, scale_group_k=32, scale_dtype=torch.float8_e8m0fnu
)
_NVFP4_WEIGHT = MegaMoEWeightSpec(
torch.int8,
values_per_element=2,
scale_group_k=16,
scale_dtype=torch.float8_e4m3fn,
)
_BF16_WEIGHT = MegaMoEWeightSpec(torch.bfloat16)
@dataclass(frozen=True)
class MegaMoEWeightFormat:
"""One routed-expert precision of the MegaMoE feature matrix."""
name: str
# DeepGEMM MMA tag; it also keys the symmetric-buffer layout.
mma_type: str
kernel: str
routed: MegaMoEWeightSpec
# Precisions a shared expert may have to run inside the same launch.
shared: tuple[MegaMoEWeightSpec, ...]
# Quantization of the dispatched hidden states; None dispatches BF16.
hidden_quant: QuantKey | None
activations: frozenset[str]
@property
def uses_nvfp4_api(self) -> bool:
"""``nvfp4_mega_moe`` has its own symmetric buffer, BLOCK_M and scales."""
return self.kernel == "nvfp4_mega_moe"
MEGA_MOE_MXFP4 = MegaMoEWeightFormat(
name="mxfp4",
mma_type="fp8xfp4",
kernel="fp8_fp4_mega_moe",
routed=_MXFP4_WEIGHT,
shared=(_MXFP8_WEIGHT,),
hidden_quant=kMxfp8Dynamic,
activations=frozenset({"swiglu", "situ"}),
)
MEGA_MOE_MXFP8 = MegaMoEWeightFormat(
name="mxfp8",
mma_type="fp8xfp8",
kernel="fp8_fp4_mega_moe",
routed=_MXFP8_WEIGHT,
shared=(_MXFP8_WEIGHT,),
hidden_quant=kMxfp8Dynamic,
activations=frozenset({"swiglu", "situ"}),
)
MEGA_MOE_NVFP4 = MegaMoEWeightFormat(
name="nvfp4",
mma_type="nvfp4",
kernel="nvfp4_mega_moe",
routed=_NVFP4_WEIGHT,
shared=(_MXFP8_WEIGHT, _BF16_WEIGHT),
hidden_quant=kNvfp4DynamicToken,
activations=frozenset({"swiglu"}),
)
MEGA_MOE_BF16 = MegaMoEWeightFormat(
name="bf16",
mma_type="bf16xbf16",
kernel="bf16_mega_moe",
routed=_BF16_WEIGHT,
shared=(_BF16_WEIGHT,),
hidden_quant=None,
activations=frozenset({"swiglu"}),
)
MEGA_MOE_WEIGHT_FORMATS = {
fmt.name: fmt
for fmt in (MEGA_MOE_MXFP4, MEGA_MOE_MXFP8, MEGA_MOE_NVFP4, MEGA_MOE_BF16)
}
def check_deep_gemm_mega_moe_supported(
device: torch.device, hidden_size: int, intermediate_size: int
) -> None:
"""Raise unless DeepGEMM MegaMoE can run on ``device`` with these shapes.
Raises:
NotImplementedError: If the device architecture is unsupported.
ValueError: If the shapes are unsupported.
"""
if not current_platform.is_device_capability_family(
100, device_id=device.index or 0
):
raise NotImplementedError("DeepGEMM MegaMoE requires SM100 GPUs.")
if hidden_size % 128 != 0 or intermediate_size % 128 != 0:
raise ValueError(
"DeepGEMM MegaMoE requires hidden and intermediate sizes "
"to be multiples of 128."
)
def _linear_weight_scale(linear: nn.Module) -> torch.Tensor | None:
scale = getattr(linear, "weight_scale", None)
return getattr(linear, "weight_scale_inv", None) if scale is None else scale
def _matches_weight_spec(linear: nn.Module, spec: MegaMoEWeightSpec) -> bool:
"""Whether a checkpoint linear can be written into a ``spec`` weight."""
weight, scale = linear.weight, _linear_weight_scale(linear)
if weight.dtype != spec.dtype or weight.dim() != 2:
return False
if spec.scale_group_k is None:
return scale is None
if scale is None or scale.dim() != 2:
return False
(rows, cols), (scale_rows, scale_cols) = weight.shape, scale.shape
scale_dtypes = {spec.scale_dtype}
if spec.scale_dtype == torch.float8_e8m0fnu:
scale_dtypes.add(torch.uint8)
return (
scale.dtype in scale_dtypes
and rows % scale_rows == 0
and cols % scale_cols == 0
and (cols // scale_cols) % spec.scale_group_k == 0
)
class DeepGemmMegaMoEExperts(nn.Module):
"""Routed experts, and optionally the shared expert, as one MegaMoE launch.
Parameters are registered in the kernel layout and every checkpoint shard is
written straight into it, so loading (and reloading) needs no post-load
transform. Whether the shared expert runs inside the launch is decided here,
before any weight loads; a fused layer owns the shared-expert weights and
the caller drops its shared MLP.
NVFP4 experts also take per-expert ``weight_scale_2`` (``w1`` and ``w3``
must match: one L1 scale per expert) and ``w2`` ``input_scale``; the kernel
quantizes every L2 input with the max of the latter over all loaded
experts. Dispatch uses per-token scales, so ``w1``/``w3`` ``input_scale``
are not loaded.
Args:
num_experts: Physical experts across EP ranks, redundant included.
activation: Activation fused into the L1 epilogue.
weight_format: Routed-expert precision, a row of the feature matrix.
shared_experts: Shared MLP to fuse if it is unsharded and in one of
the ``weight_format.shared`` precisions; ``None`` never fuses.
"""
_symm_buffers: ClassVar[dict[tuple, Any]] = {}
_synchronized_ep_groups: ClassVar[set[tuple[int, int]]] = set()
def __init__(
self,
vllm_config: VllmConfig,
*,
num_experts: int,
num_local_experts: int,
experts_start_idx: int,
top_k: int,
hidden_size: int,
intermediate_size: int,
activation: MegaMoEActivation,
weight_format: MegaMoEWeightFormat,
shared_experts: nn.Module | None = None,
num_logical_experts: int | None = None,
fast_math: bool = True,
prefix: str = "",
):
super().__init__()
if activation.name not in weight_format.activations:
raise NotImplementedError(
f"DeepGEMM MegaMoE {weight_format.name} experts support "
f"{sorted(weight_format.activations)}, not {activation.name!r}."
)
device = torch.get_default_device()
if device.type != current_platform.device_type: # meta under ipc_cache
device = torch.device(
current_platform.device_type, torch.accelerator.current_device_index()
)
check_deep_gemm_mega_moe_supported(device, hidden_size, intermediate_size)
self.prefix = prefix
self.capture_fn: Callable[[torch.Tensor], None] | None = None
self.num_experts = num_experts
self.num_local_experts = num_local_experts
self.experts_start_idx = experts_start_idx
self.experts_end_idx = experts_start_idx + num_local_experts
self.num_logical_experts = (
num_logical_experts if num_logical_experts is not None else num_experts
)
self.top_k = top_k
self.hidden_size = hidden_size
self.intermediate_size = intermediate_size
self.max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens
self.weight_format = weight_format
self.activation = activation
self.fast_math = fast_math
self._kernel_kwargs = activation.kernel_kwargs()
self.eplb_state = EplbLayerState()
# L2 input scale of every logical expert; the NVFP4 kernel takes the max.
self._w2_input_scales = torch.zeros(self.num_logical_experts, device="cpu")
self.l2_activation_scale = 1.0
# The first of each expert's w1/w3 weight_scale_2 until the other loads.
self._pending_l1_scales: dict[int, tuple[str, float]] = {}
self.num_shared_experts, self.shared_spec = self._fusable_shared_experts(
shared_experts
)
self._create_weights()
for name, param in self.named_parameters(recurse=False):
shared = name.startswith("shared_")
loader = self.shared_weight_loader if shared else self.weight_loader
set_weight_attrs(param, {"weight_loader": loader})
if not shared and name.endswith("weight_scale"):
param.quant_method = "block"
# Register in the static forward context so the custom-op wrapper
# can look up this module by name from within a torch.compile graph.
compilation_config = vllm_config.compilation_config
if prefix in compilation_config.static_forward_context:
raise ValueError(f"Duplicate layer name: {prefix}")
compilation_config.static_forward_context[prefix] = self
@property
def has_fused_shared_experts(self) -> bool:
return self.num_shared_experts > 0
def _fusable_shared_experts(
self, shared_experts: nn.Module | None
) -> tuple[int, MegaMoEWeightSpec | None]:
"""How many shared experts the launch fuses, and in which precision."""
if shared_experts is None:
return 0, None
gate_up, down = shared_experts.gate_up_proj, shared_experts.down_proj
hidden, inter = self.hidden_size, self.intermediate_size
shared_size = down.weight.shape[-1]
spec = next(
(
spec
for spec in self.weight_format.shared
if all(_matches_weight_spec(lin, spec) for lin in (gate_up, down))
),
None,
)
if (
tuple(gate_up.weight.shape) != (2 * shared_size, hidden)
or tuple(down.weight.shape) != (hidden, shared_size)
or shared_size % inter
):
reason = (
f"needs an unsharded shared MLP over hidden size {hidden} with a "
f"multiple of {inter} intermediate; got {tuple(gate_up.weight.shape)}"
f" and {tuple(down.weight.shape)}"
)
elif spec is None:
scale = _linear_weight_scale(gate_up)
supported = ", ".join(
f"{s.dtype}{' with 1x32 UE8M0 scales' if s.scale_group_k else ''}"
for s in self.weight_format.shared
)
reason = (
f"{self.weight_format.name} experts fuse {supported} shared "
f"experts; got weight {gate_up.weight.dtype} and scale "
f"{None if scale is None else (tuple(scale.shape), scale.dtype)}"
)
elif not self._kernel_fuses_shared_experts(spec):
reason = (
"the installed DeepGEMM Python API is older than the vLLM source. "
"Rebuild the vendored _deep_gemm_C extension to enable it"
)
else:
return shared_size // inter, spec
logger.warning_once("Disabling MegaMoE shared-expert fusion: %s.", reason)
return 0, None
def _kernel_fuses_shared_experts(self, spec: MegaMoEWeightSpec) -> bool:
from vllm.utils.deep_gemm import _import_deep_gemm
deep_gemm = _import_deep_gemm()
fmt = self.weight_format
buffer_api = "get_symm_buffer_for_mega_moe"
block_m_api = "get_block_m_for_mega_moe"
if fmt.uses_nvfp4_api:
buffer_api = "NVFP4SymmBuffer"
block_m_api = "get_block_m_for_nvfp4_mega_moe"
try:
api = {
*signature(getattr(deep_gemm, buffer_api)).parameters,
*signature(getattr(deep_gemm, fmt.kernel)).parameters,
}
except (AttributeError, TypeError, ValueError):
return False
# Scaled shared inputs are laid out by the runtime BLOCK_M.
return {
"num_shared_experts",
"shared_l1_weights",
"shared_l2_weights",
} <= api and (spec.scale_group_k is None or hasattr(deep_gemm, block_m_api))
def _create_weights(self) -> None:
"""Register the routed and fused shared-expert parameters.
Routed parameters are indexed by local expert on dim 0, which EPLB moves.
"""
e, h, i = self.num_local_experts, self.hidden_size, self.intermediate_size
fmt = self.weight_format
self._register_weight("w13", (e, 2 * i, h), fmt.routed)
self._register_weight("w2", (e, h, i), fmt.routed)
if fmt.routed.has_global_scale:
# w13_weight_scale_2 is the kernel's L1 alpha (per-token dispatch).
for name in ("w13_weight_scale_2", "w2_weight_scale_2", "w2_input_scale"):
scale = torch.zeros(e, dtype=torch.float32)
self.register_parameter(name, nn.Parameter(scale, requires_grad=False))
# w2 global weight scale times the L2 activation scale, per launch.
l2_alpha = torch.zeros(e, dtype=torch.float32)
self.register_buffer("_l2_alpha", l2_alpha, persistent=False)
if self.shared_spec is not None:
# The shared expert has the layout of one routed expert slice.
s = i * self.num_shared_experts
self._register_weight("shared_l1", (2 * s, h), self.shared_spec)
self._register_weight("shared_l2", (h, s), self.shared_spec)
def _register_weight(
self, name: str, shape: tuple[int, ...], spec: MegaMoEWeightSpec
) -> None:
def register(name: str, data: torch.Tensor) -> None:
self.register_parameter(name, nn.Parameter(data, requires_grad=False))
*outer, n, k = shape
register(
f"{name}_weight",
torch.zeros(*outer, n, k // spec.values_per_element, dtype=spec.dtype),
)
if spec.scale_group_k is not None:
# Block scales, MN-major, four per int32 along K.
scale_k = k // (4 * spec.scale_group_k)
scale = torch.zeros(*outer, scale_k, n, dtype=torch.int32)
register(f"{name}_weight_scale", scale.transpose(-1, -2))
def _kernel_weights(
self, name: str
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
weight = getattr(self, f"{name}_weight")
scale = getattr(self, f"{name}_weight_scale", None)
return weight if scale is None else (weight, scale)
@property
def l1_weights(self) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
return self._kernel_weights("w13")
@property
def l2_weights(self) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
return self._kernel_weights("w2")
def _map_global_expert_id(self, expert_id: int) -> list[int]:
"""Local slots that logical expert ``expert_id`` occupies on this rank."""
physical_ids: list[int] = []
for p in range(self.experts_start_idx, self.experts_end_idx):
if p % self.num_logical_experts == expert_id:
physical_ids.append(p - self.experts_start_idx)
return physical_ids
def weight_loader(
self,
param: nn.Parameter,
loaded_weight: torch.Tensor,
weight_name: str,
shard_id: str,
expert_id: int,
return_success: bool = False,
) -> bool | None:
if not loaded_weight.is_meta:
self._record_global_scale(weight_name, loaded_weight, shard_id, expert_id)
local_expert_ids = self._map_global_expert_id(expert_id)
param_prefix = "w2_" if shard_id == "w2" else "w13_"
if not local_expert_ids or param_prefix not in weight_name:
return False if return_success else None
is_global_scale = weight_name.endswith(("weight_scale_2", "input_scale"))
is_scale = weight_name.endswith("weight_scale")
for local_expert_id in local_expert_ids:
if is_global_scale:
param.data[local_expert_id].copy_(loaded_weight.reshape(()))
else:
self._load_weight(
param.data[local_expert_id], loaded_weight, shard_id, is_scale
)
return True if return_success else None
def _record_global_scale(
self, weight_name: str, scale: torch.Tensor, shard_id: str, expert_id: int
) -> None:
"""Track NVFP4 global scales the kernel takes one of per launch or expert.
Every rank sees the scales of all experts, so ranks agree and EPLB can
move experts.
"""
if weight_name.endswith("w2_input_scale"):
self._w2_input_scales[expert_id] = scale.max()
self.l2_activation_scale = float(self._w2_input_scales.max()) or 1.0
elif weight_name.endswith("w13_weight_scale_2"):
# Compare w1 and w3 of the same load, not the live (maybe stale) slot.
prev = self._pending_l1_scales.pop(expert_id, None)
value = float(scale)
if prev is None or prev[0] == shard_id:
self._pending_l1_scales[expert_id] = (shard_id, value)
elif not math.isclose(value, prev[1], rel_tol=1e-5):
raise ValueError(
f"MegaMoE NVFP4 expert {expert_id}: {prev[0]} weight_scale_2 "
f"{prev[1]} != {shard_id} weight_scale_2 {value}; the kernel "
"takes one L1 scale per expert."
)
def load_shared_expert_weight(self, name: str, loaded_weight: torch.Tensor) -> str:
"""Load one shared-expert checkpoint tensor; return the parameter name."""
if ".w2." in name or ".down_proj." in name:
param_name = "shared_l2_weight"
elif ".w1." in name or ".w3." in name:
param_name = "shared_l1_weight"
else:
raise ValueError(f"Unexpected shared-expert tensor: {name}")
if "weight_scale" in name:
param_name += "_scale"
param = getattr(self, param_name)
param.weight_loader(param, loaded_weight, name)
return f"{name.partition('shared_experts.')[0]}experts.{param_name}"
def shared_weight_loader(
self, param: nn.Parameter, loaded_weight: torch.Tensor, weight_name: str
) -> None:
shard_id = (
"w1" if ".w1." in weight_name else "w3" if ".w3." in weight_name else "w2"
)
is_scale = "weight_scale" in weight_name
self._load_weight(param.data, loaded_weight, shard_id, is_scale)
def _load_weight(
self,
dst: torch.Tensor,
src: torch.Tensor,
shard_id: str,
is_scale: bool,
) -> None:
"""Write checkpoint shard ``src`` into ``dst`` in the kernel layout.
``dst`` is one routed expert's slice or a shared-expert parameter.
"""
if shard_id not in ("w1", "w2", "w3"):
raise ValueError(f"Unsupported expert shard id: {shard_id}")
rows = dst.shape[0] // (1 if shard_id == "w2" else 2)
if is_scale:
src = self._pack_scale(src, rows, dst.shape[1] * 4)
else:
src = src.view(dst.dtype)
if tuple(src.shape) != (rows, dst.shape[1]):
raise ValueError(
f"MegaMoE {shard_id} weight shape {tuple(src.shape)} does not "
f"match the kernel shard ({rows}, {dst.shape[1]})"
)
dst, src = self._kernel_layout_shard(dst, src, shard_id, is_scale)
dst.copy_(src)
@staticmethod
def _pack_scale(sf: torch.Tensor, rows: int, cols: int) -> torch.Tensor:
"""Expand block scales to the kernel granularity, four per int32 along K."""
sf = sf.view(torch.uint8)
if rows % sf.shape[0] or cols % sf.shape[1]:
raise ValueError(
f"Scale shape {tuple(sf.shape)} does not tile ({rows}, {cols})"
)
sf = sf.repeat_interleave(rows // sf.shape[0], 0)
sf = sf.repeat_interleave(cols // sf.shape[1], 1)
return sf.contiguous().clone().view(torch.int32)
@staticmethod
def _kernel_layout_shard(
dst: torch.Tensor,
src: torch.Tensor,
shard_id: str,
is_scale: bool,
) -> tuple[torch.Tensor, torch.Tensor]:
"""Pair a checkpoint shard with its destination in the MegaMoE layout."""
k = src.shape[-1]
if shard_id == "w2":
if not is_scale:
return dst, src
# UTCCP: scale row 128*b + 32*j + i -> 128*b + 4*i + j.
return dst.view(-1, 32, 4, k).transpose(1, 2), src.reshape(-1, 4, 32, k)
half = 0 if shard_id == "w1" else 1
if not is_scale:
# Interleave gate/up by 8 rows: r -> (r // 8) * 16 + 8 * half + r % 8.
return dst.view(-1, 2, 8, k)[:, half], src.reshape(-1, 8, k)
# Interleave + UTCCP: 64b + 16j + 8q + m -> 128b + 64q + 32half + 4m + j.
dst = dst.view(-1, 2, 2, 8, 4, k)[:, :, half].permute(0, 3, 1, 2, 4)
return dst, src.reshape(-1, 4, 2, 8, k)
def get_symm_buffer(self):
from vllm.utils.deep_gemm import _import_deep_gemm
deep_gemm = _import_deep_gemm()
group = get_ep_group().device_group
key = (
id(group),
torch.accelerator.current_device_index(),
self.num_experts,
self.max_num_tokens,
self.top_k,
self.hidden_size,
self.intermediate_size,
self.num_shared_experts,
None if self.shared_spec is None else self.shared_spec.dtype,
self.weight_format.mma_type,
self.activation.name,
)
symm_buffer = self._symm_buffers.get(key)
if symm_buffer is None:
args = (
group,
self.num_experts,
self.max_num_tokens,
self.top_k,
self.hidden_size,
self.intermediate_size,
)
if self.weight_format.uses_nvfp4_api:
symm_buffer = deep_gemm.NVFP4SymmBuffer(
*args,
num_shared_experts=self.num_shared_experts,
shared_dtype=torch.float8_e4m3fn
if self.shared_spec is None
else self.shared_spec.dtype,
)
else:
symm_buffer = deep_gemm.get_symm_buffer_for_mega_moe(
*args,
num_shared_experts=self.num_shared_experts,
mma_type=self.weight_format.mma_type,
activation=self.activation.name,
)
self._symm_buffers[key] = symm_buffer
return symm_buffer
def get_block_m(self, num_tokens: int) -> int:
"""BLOCK_M the MegaMoE scheduler runs ``num_tokens`` local tokens with."""
from vllm.utils.deep_gemm import _import_deep_gemm
deep_gemm = _import_deep_gemm()
args = (
get_ep_group().world_size,
self.num_experts,
self.get_symm_buffer().num_max_tokens_per_rank,
num_tokens,
self.top_k,
)
if self.weight_format.uses_nvfp4_api:
return deep_gemm.get_block_m_for_nvfp4_mega_moe(*args)
return deep_gemm.get_block_m_for_mega_moe(*args, self.weight_format.mma_type)
def set_eplb_state(
self,
moe_layer_idx: int,
expert_load_view: torch.Tensor,
logical_to_physical_map: torch.Tensor,
logical_replica_count: torch.Tensor,
) -> None:
self.eplb_state.set_layer_state(
moe_layer_idx,
expert_load_view,
logical_to_physical_map,
logical_replica_count,
)
def get_expert_weights(self) -> list[torch.Tensor]:
def _to_eplb_view(name: str, t: torch.Tensor) -> torch.Tensor:
"""Return a (num_local_experts, -1) view with contiguous memory layout."""
assert t.shape[0] == self.num_local_experts
if t.is_contiguous():
return t.view(self.num_local_experts, -1)
elif t.dim() == 3 and t.stride(1) == 1 and t.stride(2) == t.shape[1]:
# scales have shape (E, M, N) with memory layout (E, N, M)
back = torch.transpose(t, 1, 2)
assert back.is_contiguous()
return back.view(self.num_local_experts, -1)
raise AssertionError(
f"MegaMoE EPLB {name}: non-contiguous expert tensor with "
f"unexpected layout shape={tuple(t.shape)} "
f"stride={tuple(t.stride())} dtype={t.dtype}"
)
return [
_to_eplb_view(name, param)
for name, param in self.named_parameters(recurse=False)
if not name.startswith("shared_")
]
def update_expert_map(self) -> None:
pass
@property
def layer_id(self) -> int:
return extract_layer_index(self.prefix)
def _synchronize_first_launch(self) -> None:
ep_group = get_ep_group()
key = (id(ep_group.cpu_group), torch.accelerator.current_device_index())
if key in self._synchronized_ep_groups:
return
torch.accelerator.synchronize()
torch.distributed.barrier(group=ep_group.cpu_group)
self._synchronized_ep_groups.add(key)
def forward(
self,
hidden_states: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
shared_experts: nn.Module | None = None,
) -> torch.Tensor:
"""Run the routed and shared experts on this rank's tokens.
Args:
hidden_states: BF16 ``[num_tokens, hidden_size]`` activations.
topk_weights: ``[num_tokens, top_k]`` routing weights.
topk_ids: ``[num_tokens, top_k]`` logical expert ids.
shared_experts: Unfused shared MLP, run on ``hidden_states`` and
added to the output. Must be ``None`` when this layer fused the
shared expert at construction.
Returns:
BF16 ``[num_tokens, hidden_size]`` routed plus shared expert output.
"""
if shared_experts is not None and self.has_fused_shared_experts:
raise ValueError(
f"{self.prefix} fuses the shared expert into MegaMoE; "
"pass shared_experts=None."
)
num_tokens = hidden_states.shape[0]
if num_tokens > self.max_num_tokens:
raise ValueError(
f"MegaMoE got {num_tokens} tokens, but the symmetric buffer was "
f"sized for {self.max_num_tokens}."
)
self._synchronize_first_launch()
is_padding = None
if envs.VLLM_MOE_SKIP_PADDING and is_forward_context_available():
is_padding = get_forward_context().is_padding
if is_padding is not None:
is_padding = is_padding[:num_tokens]
if self.capture_fn is not None:
self.capture_fn(topk_ids)
# EPLB: map logical expert IDs to physical replicas and record load.
eplb_state = self.eplb_state
if eplb_state.logical_to_physical_map is not None:
assert eplb_state.expert_load_view is not None
assert eplb_state.logical_replica_count is not None
assert eplb_state.should_record_tensor is not None
if is_padding is not None:
topk_ids = torch.where(is_padding.unsqueeze(1), -1, topk_ids)
topk_ids = eplb_map_to_physical_and_record(
topk_ids=topk_ids,
expert_load_view=eplb_state.expert_load_view,
logical_to_physical_map=eplb_state.logical_to_physical_map,
logical_replica_count=eplb_state.logical_replica_count,
record_enabled=eplb_state.should_record_tensor,
num_unpadded_tokens=eplb_state.num_unpadded_tokens_tensors[
dbo_current_ubatch_id()
]
if eplb_state.num_unpadded_tokens_tensors is not None
else None,
)
from vllm.utils.deep_gemm import _import_deep_gemm
deep_gemm = _import_deep_gemm()
fmt = self.weight_format
symm_buffer = self.get_symm_buffer()
shared_x_sf = None
shared_block_m = None
if self.shared_spec is not None and self.shared_spec.scale_group_k:
shared_x_sf = symm_buffer.shared_l1_acts_sf
shared_block_m = self.get_block_m(num_tokens)
kwargs = dict(self._kernel_kwargs)
if fmt.uses_nvfp4_api:
l2_activation_scale = self.l2_activation_scale
prepare_megamoe_nvfp4_inputs(
hidden_states,
topk_weights,
topk_ids,
symm_buffer.x[:num_tokens],
symm_buffer.x_sf[:num_tokens],
symm_buffer.x_scales[:num_tokens],
symm_buffer.topk_idx[:num_tokens],
symm_buffer.topk_weights[:num_tokens],
is_padding=is_padding,
shared_x=symm_buffer.shared_l1_acts[:num_tokens]
if self.has_fused_shared_experts
else None,
shared_x_sf=shared_x_sf,
shared_block_m=shared_block_m,
w2_global_scale=self.w2_weight_scale_2,
l2_alpha=self._l2_alpha,
l2_activation_scale=l2_activation_scale,
)
kwargs.update(
l1_alpha=self.w13_weight_scale_2,
l2_alpha=self._l2_alpha,
l2_activation_scale=l2_activation_scale,
use_x_scales=True,
)
else:
quantized = fmt.hidden_quant is not None
prepare_megamoe_inputs(
hidden_states,
topk_weights,
topk_ids,
symm_buffer.x[:num_tokens],
symm_buffer.x_sf[:num_tokens] if quantized else None,
symm_buffer.topk_idx[:num_tokens],
symm_buffer.topk_weights[:num_tokens],
is_padding=is_padding,
shared_x_sf=shared_x_sf,
shared_block_m=shared_block_m,
hidden_quant=fmt.hidden_quant,
)
if self.has_fused_shared_experts:
kwargs["shared_l1_weights"] = self._kernel_weights("shared_l1")
kwargs["shared_l2_weights"] = self._kernel_weights("shared_l2")
y = torch.empty_like(hidden_states, dtype=torch.bfloat16)
getattr(deep_gemm, fmt.kernel)(
y,
self.l1_weights,
self.l2_weights,
symm_buffer,
fast_math=self.fast_math,
**kwargs,
)
if shared_experts is not None:
y += shared_experts(hidden_states)
return y
DeepGemmMegaMoEExperts.weight_loader.supports_moe_loading = True # type: ignore[attr-defined]
vllm/model_executor/warmup/jit_warmup_cutedsl_helper.py
-79
View File
@@ -1,79 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from abc import abstractmethod
from collections.abc import Callable, Mapping
from typing import Any, ClassVar, Generic, TypeAlias, TypeVar
from vllm.model_executor.warmup.jit_warmup import (
VllmJitKernel,
)
DEFAULT_CUTEDSL_COMPILE_OPTIONS = "--enable-tvm-ffi"
CompileKeyT = TypeVar("CompileKeyT")
CuTeDSLLaunchSpec: TypeAlias = (
tuple[CompileKeyT, tuple[Any, ...]]
| tuple[CompileKeyT, tuple[Any, ...], Any]
| tuple[
CompileKeyT,
tuple[Any, ...],
Any,
Callable[[], Any],
]
)
def cutedsl_fake_stream(*, use_tvm_ffi_env_stream: bool = True) -> Any:
import cutlass.cute as cute
return cute.runtime.make_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream)
def compile_cutedsl(
entry: Callable[..., Any],
*args: Any,
options: str = DEFAULT_CUTEDSL_COMPILE_OPTIONS,
use_tvm_ffi_env_stream: bool = True,
) -> Any:
import cutlass.cute as cute
return cute.compile(
entry,
*args,
cutedsl_fake_stream(use_tvm_ffi_env_stream=use_tvm_ffi_env_stream),
options=options,
)
class VllmCuTeDSLJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
"""CuTeDSL owner whose compiled executor is shared by warmup and runtime."""
kernel: ClassVar[Any]
bind_launch_inputs = False
@abstractmethod
def warmup_inputs(self, compile_key: CompileKeyT) -> tuple[Any, ...]:
"""Return fake arguments that compile one executor specialization."""
raise NotImplementedError
def compile(self, compile_key: CompileKeyT) -> None:
if compile_key in self._compiled_cache:
return
self._compiled_cache[compile_key] = compile_cutedsl(
self.kernel(compile_key),
*self.warmup_inputs(compile_key),
)
def launch(
self,
launch_spec: CuTeDSLLaunchSpec[CompileKeyT],
_inputs: Mapping[str, Any],
) -> Any:
# (compile_key, args), optionally followed by output and epilogue.
compile_key, launch_args = launch_spec[:2]
executor = self._get_or_compile(compile_key)
result = executor(*launch_args)
if len(launch_spec) == 4:
return launch_spec[3]()
return launch_spec[2] if len(launch_spec) == 3 else result
vllm/model_executor/warmup/jit_warmup_tilelang_helper.py
-171
View File
@@ -1,171 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Compile-only helpers for TileLang JIT warmup."""
from __future__ import annotations
import logging
from abc import abstractmethod
from collections.abc import Callable, Iterator, Mapping, MutableMapping
from contextlib import contextmanager
from dataclasses import dataclass, field
from functools import wraps
from typing import Any, ClassVar, Generic, TypeAlias, TypeVar, cast
import torch
from vllm.model_executor.warmup.jit_warmup import VllmJitKernel
CompileKeyT = TypeVar("CompileKeyT")
# (kernel_args, launch_args), optionally followed by keywords and a result.
TileLangLaunchSpec: TypeAlias = (
tuple[tuple[Any, ...], tuple[Any, ...]]
| tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None]
| tuple[tuple[Any, ...], tuple[Any, ...], Mapping[str, Any] | None, Any]
)
@contextmanager
def _quiet_tilelang_warmup_logs() -> Iterator[None]:
loggers = [
logging.getLogger("tilelang.cache.kernel_cache"),
logging.getLogger("tilelang.jit.kernel"),
]
previous_levels = [logger.level for logger in loggers]
try:
for logger in loggers:
logger.setLevel(logging.ERROR)
yield
finally:
for logger, level in zip(loggers, previous_levels):
logger.setLevel(level)
@dataclass(frozen=True)
class TileLangWarmupTensor:
"""Minimal tensor-like object accepted by TileLang compile().
TileLang builds its cache key and TIR from tensor dtype, shape and stride.
This object deliberately has no storage, so compile-only warmup does not
allocate GPU memory and does not launch the kernel.
"""
dtype: torch.dtype
shape: tuple[int, ...] = (1,)
strides: tuple[int, ...] | None = field(default=None)
def stride(self) -> tuple[int, ...]:
if self.strides is not None:
return self.strides
strides: list[int] = []
stride = 1
for size in reversed(self.shape):
strides.append(stride)
stride *= size
return tuple(reversed(strides))
def new_empty(
self,
shape: tuple[int, ...],
*,
dtype: torch.dtype | None = None,
device: Any = None,
) -> TileLangWarmupTensor:
return TileLangWarmupTensor(dtype or self.dtype, shape)
def make_tilelang_warmup_tensor(
dtype: torch.dtype,
*shape: int,
strides: tuple[int, ...] | None = None,
) -> TileLangWarmupTensor:
return TileLangWarmupTensor(dtype=dtype, shape=tuple(shape), strides=strides)
def _tilelang_call_kwargs(kwargs: dict[str, Any]) -> dict[str, Any]:
call_kwargs = dict(kwargs)
tune_params = call_kwargs.pop("__tune_params", {})
if isinstance(tune_params, dict):
call_kwargs.update(tune_params)
return call_kwargs
def compile_tilelang(jit_impl: Any, *args: Any, **kwargs: Any) -> None:
"""Compile one TileLang specialization and populate its call cache.
TileLang's ``compile()`` materializes the kernel without launching it.
We also store the compiled kernel in ``_kernel_cache`` using the same
parsed key as ``__call__`` so runtime does not report a cache miss for an
already materialized specialization.
"""
with _quiet_tilelang_warmup_logs():
compiled = jit_impl.compile(*args, **kwargs)
func = getattr(jit_impl, "func", None)
parse_args = getattr(func, "parse_args", None)
cache = getattr(jit_impl, "_kernel_cache", None)
if not callable(parse_args) or not isinstance(cache, MutableMapping):
return
key, _ = parse_args(*args, **_tilelang_call_kwargs(kwargs))
cache[key] = compiled
class VllmTileLangJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]):
"""TileLang owner whose runtime launch specification is reused for warmup."""
kernel: ClassVar[Any]
_warming_key: CompileKeyT | None = None
@abstractmethod
def warmup_inputs(self, compile_key: CompileKeyT) -> dict[str, Any]:
"""Return runtime-shaped inputs that reproduce one compile key."""
raise NotImplementedError
def compile(self, compile_key: CompileKeyT) -> None:
self._warming_key = compile_key
try:
cast(Callable[..., Any], self)(**self.warmup_inputs(compile_key))
finally:
self._warming_key = None
def launch(
self,
jit_impl: Any,
args: tuple[Any, ...],
kwargs: Mapping[str, Any] | None = None,
) -> Any:
call_kwargs = dict(kwargs or {})
if self._warming_key is not None:
return compile_tilelang(jit_impl, *args, **call_kwargs)
return jit_impl(*args, **call_kwargs)
def kernel_launcher(
call_fn: Callable[..., TileLangLaunchSpec],
) -> Callable[..., Any]:
"""Launch TileLang from declarative kernel and runtime argument tuples."""
@wraps(call_fn)
def wrapper(
self: VllmTileLangJitKernel[Any],
*args: Any,
**kwargs: Any,
) -> Any:
launch_spec = call_fn(self, *args, **kwargs)
kernel_args, launch_args = launch_spec[:2]
if self._warming_key is not None and kernel_args not in (
(),
(self._warming_key,),
):
raise RuntimeError(
"TileLang warmup inputs produced a different compile key: "
f"expected {self._warming_key!r}, got kernel arguments "
f"{kernel_args!r}"
)
launch_kwargs = launch_spec[2] if len(launch_spec) > 2 else None
output = launch_spec[3] if len(launch_spec) > 3 else None
result = self.launch(self.kernel(*kernel_args), launch_args, launch_kwargs)
return output if len(launch_spec) > 3 else result
return wrapper
vllm/models/common/ops/fused_allreduce_rms_norm.py
-63
View File
@@ -1,63 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Fused all-reduce + residual-add + RMSNorm for eager model paths.
This recovers a fusion that vLLM's torch.compile passes would normally do but
that doesn't fire for models running eager (or under a breakable CUDA graph).
"""
import torch
from vllm.distributed import (
get_tensor_model_parallel_world_size,
tensor_model_parallel_all_reduce,
)
from vllm.model_executor.layers.fused_allreduce_gemma_rms_norm import (
_AR_RESIDUAL_RMS_NORM,
_can_use_flashinfer,
flashinfer_trtllm_fused_allreduce_norm,
)
from vllm.model_executor.layers.layernorm import RMSNorm
def fused_allreduce_rms_norm(
hidden_states: torch.Tensor,
residual: torch.Tensor,
norm: RMSNorm,
) -> tuple[torch.Tensor, torch.Tensor]:
"""All-reduce + add residual + (standard) RMSNorm, fused via flashinfer.
``hidden_states`` is the per-rank *partial* output of a row-parallel linear
run with ``reduce_results=False``; ``norm`` is the RMSNorm applied right
after. Returns ``(normed_output, new_residual)``, equivalent to
``norm(all_reduce(hidden_states), residual)``. Falls back to an explicit
all-reduce + RMSNorm when the flashinfer fast path is unavailable.
"""
tp_size = get_tensor_model_parallel_world_size()
if tp_size == 1:
return norm(hidden_states, residual)
if flashinfer_trtllm_fused_allreduce_norm is not None:
ok, max_token_num = _can_use_flashinfer(hidden_states, tp_size)
if ok:
norm_out = torch.empty_like(hidden_states)
# With norm_out provided, the kernel writes the new residual
# (all_reduce(hidden_states) + residual) into the hidden_states
# buffer and the normalized result into norm_out.
flashinfer_trtllm_fused_allreduce_norm(
allreduce_in=hidden_states,
residual=residual,
rms_gamma=norm.weight,
rms_eps=norm.variance_epsilon,
world_size=tp_size,
weight_bias=0.0, # standard RMSNorm (Gemma would use 1.0)
launch_with_pdl=True,
fp32_acc=True,
max_token_num=max_token_num,
pattern_code=_AR_RESIDUAL_RMS_NORM,
norm_out=norm_out,
)
return norm_out, hidden_states
reduced = tensor_model_parallel_all_reduce(hidden_states)
return norm(reduced, residual)
vllm/multimodal/video.py
-1
View File
@@ -195,7 +195,6 @@ class VideoLoader:
VIDEO_LOADER_REGISTRY = VideoLoaderRegistry()
VIDEO_LOADER_REGISTRY.register_gpu_codec("deepstream")
@VIDEO_LOADER_REGISTRY.register("opencv")
vllm/multimodal/video_decoders/__init__.py
+1 -10
View File
@@ -12,24 +12,15 @@ from .base import (
check_frame_pixel_limit,
)
VideoDecoderBackend = Literal["opencv", "torchcodec", "pynvvideocodec", "deepstream"]
VideoDecoderBackend = Literal["opencv", "pynvvideocodec"]
_BACKEND_OPTION_DEFAULTS: dict[str, dict[str, Any]] = {
"opencv": {},
"torchcodec": {
"num_ffmpeg_threads": 0,
"seek_mode": "exact",
"device": "cpu",
},
PYNVVIDEOCODEC_VIDEO_BACKEND: {
"hw_decoders": PYNVVIDEOCODEC_DEFAULT_HW_DECODERS,
"max_width": None,
"max_height": None,
},
"deepstream": {
"pool_size": None,
"timeout_sec": 120.0,
},
}
vllm/multimodal/video_decoders/deepstream.py
-151
View File
@@ -1,151 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import threading
from typing import Any, ClassVar
import numpy.typing as npt
from vllm.logger import init_logger
from .base import VideoSourceMetadata, VideoTargetMetadata, check_frame_pixel_limit
logger = init_logger(__name__)
def decode_deepstream(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
*,
pool_size: int | None = None,
timeout_sec: float = 120.0,
) -> tuple[npt.NDArray, VideoSourceMetadata, list[int], list[int]]:
from nvidia.deepstream_videodecode import probe_metadata
total_frames, original_fps, duration, _width, _height, codec = probe_metadata(data)
check_frame_pixel_limit(_width, _height)
source = loader_cls._prepare_source(
VideoSourceMetadata(total_frames, original_fps, duration)
)
frame_idx = loader_cls.compute_frames_index_to_sample(
source=source, target=target, **sampling_kwargs
)
frames, valid = DeepStreamVideoBackendMixin.decode_indices(
data,
frame_idx,
source,
codec=codec,
pool_size=pool_size,
timeout_sec=timeout_sec,
)
return frames, source, frame_idx, valid
class DeepStreamVideoBackendMixin:
"""NVIDIA DeepStream (NVDEC) GPU-decode codec utilities.
Decoding runs on a shared pool of daemon threads inside one CUDA
context (see the ``nvidia-deepstream-videodecode-cu13`` package). The
container bytes are pushed into an ``appsrc`` GStreamer pipeline, so no
local file path is required — HTTP and base64 sources decode identically
to local files.
Like the OpenCV mixin, this provides only the codec layer.
Frame *selection* lives in the loader's
``compute_frames_index_to_sample`` and arrives here as an explicit
list of frame indices.
"""
# Process-wide lazy decode pool, shared across all DeepStream backends.
_pool: ClassVar[Any] = None
_pool_lock: ClassVar[Any] = None
@classmethod
def _get_pool(cls, pool_size: int | None = None):
"""Lazy-initialize the shared decode pool on first use.
``pool_size`` (number of decode worker threads) comes from
``--media-io-kwargs`` (``{"video": {"pool_size": N}}``); when unset it
defaults to the existing ``VLLM_MEDIA_LOADING_THREAD_COUNT`` so no
DeepStream-specific env var is needed. The pool is a process-wide
singleton, so the first decode's value wins.
"""
if cls._pool is not None:
return cls._pool
if cls._pool_lock is None:
cls._pool_lock = threading.Lock()
with cls._pool_lock:
if cls._pool is not None:
return cls._pool
import os
from nvidia.deepstream_videodecode import DecodePool
if pool_size is None:
pool_size = int(os.environ.get("VLLM_MEDIA_LOADING_THREAD_COUNT", 8))
pool_size = max(1, min(int(pool_size), 16))
logger.info(
"[DeepStream] initializing decode pool with %d workers",
pool_size,
)
cls._pool = DecodePool(num_workers=pool_size)
return cls._pool
@classmethod
def decode_indices(
cls,
data: bytes,
frame_indices: list[int],
source: VideoSourceMetadata,
codec: str = "",
pool_size: int | None = None,
timeout_sec: float = 120.0,
) -> tuple[npt.NDArray, list[int]]:
"""Decode the requested frame indices from raw container bytes.
The whole stream is decoded; the pool keeps exactly the frames whose
decode-order index is in ``frame_indices`` (1:1, frame-exact) and
sends EOS once the last one is matched.
``codec`` (e.g. ``"h264"``/``"hevc"``) lets the pool keep its NVDEC
session warm across same-codec streams and rebuild only on a codec
change. Frames are returned as a CPU NHWC uint8 array so the
upstream multimodal parser sees the same shape as the other
backends.
"""
if not frame_indices:
raise ValueError("DeepStream backend received no frame indices")
result = cls._get_pool(pool_size).decode(
data,
target_indices=frame_indices,
codec=codec,
max_frames=len(frame_indices),
timeout_sec=timeout_sec,
)
if result.error:
raise ValueError(f"DeepStream decode failed: {result.error}")
if result.frames is None or result.n_kept == 0:
raise ValueError("DeepStream decode produced no frames")
valid = frame_indices[: result.n_kept]
# GPU -> CPU NHWC uint8 at the codec boundary (one PCIe copy); keeps
# the array shape identical to the OpenCV backend. Copy into
# PINNED host memory (reused across calls by PyTorch's pinned caching
# allocator) so the D2H runs at full PCIe bandwidth (~13 GB/s) rather
# than the ~1 GB/s pageable path that plain ``.cpu()`` takes — ~12x
# faster for a 1080p x8 frame batch (~46ms -> ~4ms). ``numpy()`` keeps
# the pinned tensor alive via the array's base.
import torch
gpu = result.frames
if gpu.is_cuda:
host = torch.empty(gpu.shape, dtype=gpu.dtype, pin_memory=True)
host.copy_(gpu, non_blocking=True)
torch.cuda.current_stream().synchronize()
arr = host.numpy()
else:
arr = gpu.numpy()
return arr, valid
vllm/multimodal/video_decoders/torchcodec.py
-133
View File
@@ -1,133 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from typing import Literal
import numpy as np
import numpy.typing as npt
import torch
from vllm.logger import init_logger
from vllm.platforms import current_platform
from vllm.utils.import_utils import PlaceholderModule, check_torchcodec_available
from .base import (
VideoSourceMetadata,
VideoTargetMetadata,
check_frame_pixel_limit,
)
try:
from torchcodec.decoders import VideoDecoder
except (ImportError, RuntimeError):
VideoDecoder = PlaceholderModule("torchcodec").placeholder_attr( # type: ignore[assignment]
"decoders.VideoDecoder"
)
logger = init_logger(__name__)
def decode_torchcodec(
loader_cls,
data: bytes,
target: VideoTargetMetadata,
sampling_kwargs: dict,
*,
num_ffmpeg_threads: int = 0,
seek_mode: Literal["exact", "approximate"] = "exact",
device: str = "cpu",
) -> tuple[npt.NDArray | torch.Tensor, VideoSourceMetadata, list[int], list[int]]:
check_torchcodec_available()
decoder = TorchCodecVideoBackendMixin.make_torchcodec_decoder(
data,
num_ffmpeg_threads=num_ffmpeg_threads,
seek_mode=seek_mode,
device=device,
)
check_frame_pixel_limit(
decoder.metadata.width or 0,
decoder.metadata.height or 0,
)
source = loader_cls._prepare_source(
TorchCodecVideoBackendMixin.get_torchcodec_metadata(decoder)
)
frame_idx = loader_cls.compute_frames_index_to_sample(
source=source, target=target, **sampling_kwargs
)
frames, valid = TorchCodecVideoBackendMixin.decode_torchcodec_frames(
decoder, frame_idx, device=device
)
return frames, source, frame_idx, valid
class TorchCodecVideoBackendMixin:
"""TorchCodec (FFmpeg-backed, PyTorch-native) codec utilities.
Builds a :class:`~torchcodec.decoders.VideoDecoder` over the in-memory
bytes and extracts the sampled indices with a single batched
``get_frames_at`` call, while releasing the GIL during decode.
"""
@staticmethod
def make_torchcodec_decoder(
data: bytes,
*,
num_ffmpeg_threads: int = 0,
seek_mode: Literal["exact", "approximate"] = "exact",
device: str = "cpu",
) -> "VideoDecoder":
torch_device = torch.device(device)
if torch_device.type == "cuda" and not current_platform.is_cuda():
raise ValueError(
f"torchcodec video decoding on device {device!r} requires "
"a CUDA-capable platform."
)
elif torch_device.type not in ("cpu", "cuda"):
raise ValueError(
f"torchcodec video decoding only supports 'cpu' and 'cuda' "
f"devices, got {device!r}."
)
# NHWC matches the (num_frames, H, W, 3) uint8 RGB layout the rest
# of the pipeline expects, avoiding a transpose.
return VideoDecoder(
data,
dimension_order="NHWC",
num_ffmpeg_threads=num_ffmpeg_threads,
seek_mode=seek_mode,
device=device,
)
@staticmethod
def get_torchcodec_metadata(decoder: "VideoDecoder") -> VideoSourceMetadata:
md = decoder.metadata
total_frames = md.num_frames or 0
fps = float(md.average_fps) if md.average_fps else 0.0
duration = float(md.duration_seconds) if md.duration_seconds else 0.0
if total_frames == 0 and duration > 0 and fps > 0:
total_frames = int(duration * fps)
return VideoSourceMetadata(total_frames, fps, duration)
@staticmethod
def decode_torchcodec_frames(
decoder: "VideoDecoder",
frame_indices: list[int],
*,
device: str = "cpu",
) -> tuple[npt.NDArray | torch.Tensor, list[int]]:
"""Decode the requested indices in one batched, index-exact call."""
if not frame_indices:
return np.empty((0,), dtype=np.uint8), []
# Note: torchcodec releases the GIL for the entire call
batch = decoder.get_frames_at(frame_indices)
frames = batch.data
if (
torch.device(device).type != "cpu"
and (status := getattr(decoder, "cpu_fallback", None)) is not None
and bool(status)
):
logger.warning_once(
"torchcodec could not use NVDEC for this video and "
"decoded on CPU instead; check codec support and "
"libnvcuvid."
)
return frames, list(frame_indices)
vllm/profiler/layerwise_profile.py
-400
View File
@@ -1,400 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import copy
from collections import defaultdict
from collections.abc import Callable
from dataclasses import asdict, dataclass, field
from typing import Any, Generic, TypeAlias, TypeVar
from torch._C._autograd import DeviceType, _KinetoEvent, _ProfilerResult
from torch._C._profiler import _EventType, _ExperimentalConfig, _ProfilerEvent
from torch.autograd.profiler import FunctionEvent
from torch.profiler import ProfilerActivity, profile
from vllm.profiler.utils import (
TablePrinter,
event_has_module,
event_is_torch_op,
event_module_repr,
event_torch_op_stack_trace,
indent_string,
)
from vllm.utils.import_utils import PlaceholderModule
try:
import pandas as pd
except ImportError:
pd = PlaceholderModule("pandas")
@dataclass
class _ModuleTreeNode:
event: _ProfilerEvent
parent: "_ModuleTreeNode | None" = None
children: list["_ModuleTreeNode"] = field(default_factory=list)
trace: str = ""
@property
def is_leaf(self):
return self.event.children is None or len(self.event.children) == 0
@property
def is_torch_op(self):
return event_is_torch_op(self.event)
@property
def is_cuda(self):
return (
self.event.tag == _EventType.Kineto
and self.event.typed[1].device_type == DeviceType.CUDA
)
@dataclass
class SummaryStatsEntry:
name: str
cuda_time_us: float
pct_cuda_time: float
invocations: int
@dataclass
class ModelStatsEntry:
name: str
cpu_time_us: float
cuda_time_us: float
pct_cuda_time: float
trace: str
StatsEntry: TypeAlias = ModelStatsEntry | SummaryStatsEntry
StatsEntryT = TypeVar("StatsEntryT", bound=StatsEntry)
@dataclass
class _StatsTreeNode(Generic[StatsEntryT]):
entry: StatsEntryT
children: list["_StatsTreeNode[StatsEntryT]"] = field(default_factory=list)
parent: "_StatsTreeNode[StatsEntryT] | None" = None
@dataclass
class LayerwiseProfileResults(profile):
_kineto_results: _ProfilerResult
_kineto_event_correlation_map: dict[int, list[_KinetoEvent]] = field(init=False)
_event_correlation_map: dict[int, list[FunctionEvent]] = field(init=False)
_module_tree: list[_ModuleTreeNode] = field(init=False)
_model_stats_tree: list[_StatsTreeNode[ModelStatsEntry]] = field(init=False)
_summary_stats_tree: list[_StatsTreeNode[SummaryStatsEntry]] = field(init=False)
# profile metadata
num_running_seqs: int | None = None
def __post_init__(self):
self._build_correlation_map()
self._build_module_tree()
self._build_stats_trees()
def print_model_table(self, column_widths: dict[str, int] | None = None):
_column_widths = dict(
name=60, cpu_time_us=12, cuda_time_us=12, pct_cuda_time=12, trace=60
)
if column_widths:
_column_widths.update(**column_widths)
filtered_model_table = [
(depth, row)
for depth, row in self._flatten_stats_tree(self._model_stats_tree)
if row.cuda_time_us > 0 or row.cpu_time_us > 0
]
TablePrinter(ModelStatsEntry, _column_widths).print_table(
self._indent_row_names_based_on_depth(
filtered_model_table,
indent_style=lambda indent: "|" + "-" * indent + " ",
)
)
def print_summary_table(self, column_widths: dict[str, int] | None = None):
_column_widths = dict(
name=80, cuda_time_us=12, pct_cuda_time=12, invocations=15
)
if column_widths:
_column_widths.update(**column_widths)
filtered_summary_table = [
(depth, row)
for depth, row in self._flatten_stats_tree(self._summary_stats_tree)
if row.cuda_time_us > 0
]
TablePrinter(SummaryStatsEntry, _column_widths).print_table(
self._indent_row_names_based_on_depth(
filtered_summary_table,
indent_style=lambda indent: "|" + "-" * indent + " ",
)
)
def export_model_stats_table_csv(self, filename: str):
df = pd.DataFrame(
[asdict(row) for _, row in self._flatten_stats_tree(self._model_stats_tree)]
)
df.to_csv(filename)
def export_summary_stats_table_csv(self, filename: str):
df = pd.DataFrame(
[
asdict(row)
for _, row in self._flatten_stats_tree(self._summary_stats_tree)
]
)
df.to_csv(filename)
def convert_stats_to_dict(self) -> dict[str, Any]:
return {
"metadata": {"num_running_seqs": self.num_running_seqs},
"summary_stats": self._convert_stats_tree_to_dict(self._summary_stats_tree),
"model_stats": self._convert_stats_tree_to_dict(self._model_stats_tree),
}
@staticmethod
def _indent_row_names_based_on_depth(
depths_rows: list[tuple[int, StatsEntryT]],
indent_style: Callable[[int], str] | str = " ",
):
indented_rows: list[StatsEntryT] = []
for depth, row in depths_rows:
if row.cuda_time_us == 0:
continue
indented_row: StatsEntryT = copy.deepcopy(row)
indented_row.name = indent_string(indented_row.name, depth, indent_style)
indented_rows.append(indented_row)
return indented_rows
def _build_correlation_map(self):
self._kineto_event_correlation_map = defaultdict(list)
for event in self._kineto_results.events():
self._kineto_event_correlation_map[event.correlation_id()].append(event)
def _build_module_tree(self):
self._module_tree = []
event_tree = self._kineto_results.experimental_event_tree()
def _df_traversal(
event: _ProfilerEvent, curr_node: _ModuleTreeNode | None = None
):
# For the tensor parallel case for now only look at task 1
if event.start_tid != 1:
return
if event_has_module(event):
node = _ModuleTreeNode(event=event, parent=curr_node)
if curr_node:
curr_node.children.append(node)
else:
self._module_tree.append(node)
curr_node = node
is_leaf = event.children is None or len(event.children) == 0
if is_leaf and curr_node:
node = _ModuleTreeNode(
event=event,
parent=curr_node,
trace=event_torch_op_stack_trace(
event, until=lambda x: event_has_module(x)
),
)
curr_node.children.append(node)
curr_node = node
for child in event.children:
_df_traversal(child, curr_node)
for root in event_tree:
_df_traversal(root)
def _get_kineto_gpu_event(self, node: _ModuleTreeNode):
if node.event.tag != _EventType.Kineto:
return None
correlated_kineto_events = self._kineto_event_correlation_map.get(
node.event.correlation_id, []
)
iterator = (
x
for x in correlated_kineto_events
if x.device_type() == DeviceType.CUDA and x.name() == node.event.name
)
return next(iterator, None)
def _cumulative_cuda_time(self, node: _ModuleTreeNode):
"""Return cuda time in microseconds."""
def _cumulative_cuda_time_recursive(node: _ModuleTreeNode):
if node.is_leaf and (gpu_kineto_event := self._get_kineto_gpu_event(node)):
return gpu_kineto_event.duration_ns() / 1000.0
else:
cumulative_cuda_time = 0
for child in node.children:
cumulative_cuda_time += _cumulative_cuda_time_recursive(child)
return cumulative_cuda_time
return _cumulative_cuda_time_recursive(node)
def _total_cuda_time(self):
return sum([self._cumulative_cuda_time(root) for root in self._module_tree])
def _build_stats_trees(self):
summary_dict: dict[tuple[str, ...], _StatsTreeNode[SummaryStatsEntry]] = {}
total_cuda_time = self._total_cuda_time()
def pct_cuda_time(cuda_time_us):
return (cuda_time_us / total_cuda_time) * 100
def build_summary_stats_tree_df(
node: _ModuleTreeNode,
parent: _StatsTreeNode[SummaryStatsEntry] | None = None,
summary_trace: tuple[str, ...] = (),
) -> _StatsTreeNode[SummaryStatsEntry] | None:
if event_has_module(node.event):
name = event_module_repr(node.event)
cuda_time_us = self._cumulative_cuda_time(node)
elif gpu_kineto_event := self._get_kineto_gpu_event(node):
name = gpu_kineto_event.name()
cuda_time_us = gpu_kineto_event.duration_ns() / 1000.0
else:
return None
summary_trace = summary_trace + (name,)
if summary_trace in summary_dict:
entry = summary_dict[summary_trace].entry
entry.cuda_time_us += cuda_time_us
entry.invocations += 1
entry.pct_cuda_time = pct_cuda_time(entry.cuda_time_us)
else:
new_node = _StatsTreeNode(
entry=SummaryStatsEntry(
name=name,
cuda_time_us=cuda_time_us,
pct_cuda_time=pct_cuda_time(cuda_time_us),
invocations=1,
),
parent=parent,
)
if parent:
parent.children.append(new_node)
summary_dict[summary_trace] = new_node
for child in node.children:
build_summary_stats_tree_df(
child, summary_dict[summary_trace], summary_trace
)
return summary_dict[summary_trace]
self._summary_stats_tree = []
for root in self._module_tree:
summary_node = build_summary_stats_tree_df(root)
if summary_node is not None:
self._summary_stats_tree.append(summary_node)
def build_model_stats_tree_df(
node: _ModuleTreeNode,
parent: _StatsTreeNode[ModelStatsEntry] | None = None,
) -> _StatsTreeNode[ModelStatsEntry] | None:
if event_has_module(
node.event,
):
name = event_module_repr(node.event)
cuda_time_us = self._cumulative_cuda_time(node)
cpu_time_us = node.event.duration_time_ns / 1000
trace = ""
elif gpu_kineto_event := self._get_kineto_gpu_event(node):
name = gpu_kineto_event.name()
cuda_time_us = gpu_kineto_event.duration_ns() / 1000.0
cpu_time_us = 0
trace = node.trace
else:
return None
new_node = _StatsTreeNode(
entry=ModelStatsEntry(
name=name,
cpu_time_us=cpu_time_us,
cuda_time_us=cuda_time_us,
pct_cuda_time=pct_cuda_time(cuda_time_us),
trace=trace,
),
parent=parent,
)
if parent:
parent.children.append(new_node)
for child in node.children:
build_model_stats_tree_df(child, new_node)
return new_node
self._model_stats_tree = []
for root in self._module_tree:
model_node = build_model_stats_tree_df(root)
if model_node is not None:
self._model_stats_tree.append(model_node)
def _flatten_stats_tree(
self, tree: list[_StatsTreeNode[StatsEntryT]]
) -> list[tuple[int, StatsEntryT]]:
entries: list[tuple[int, StatsEntryT]] = []
def df_traversal(node: _StatsTreeNode[StatsEntryT], depth: int = 0):
entries.append((depth, node.entry))
for child in node.children:
df_traversal(child, depth=depth + 1)
for root in tree:
df_traversal(root)
return entries
def _convert_stats_tree_to_dict(
self, tree: list[_StatsTreeNode[StatsEntryT]]
) -> list[dict[str, Any]]:
root_dicts: list[dict[str, Any]] = []
def df_traversal(
node: _StatsTreeNode[StatsEntryT], curr_json_list: list[dict[str, Any]]
):
curr_json_list.append({"entry": asdict(node.entry), "children": []})
for child in node.children:
df_traversal(child, curr_json_list[-1]["children"])
for root in tree:
df_traversal(root, root_dicts)
return root_dicts
class layerwise_profile(profile):
def __init__(self, num_running_seqs: int | None = None):
"""Layerwise profile constructor.
Args:
num_running_seqs (Optional[int], optional): When given,
num_running_seqs will be passed to LayerProfileResults
for metadata update. Defaults to None.
"""
super().__init__(
activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA],
record_shapes=True,
with_stack=True,
with_modules=True,
experimental_config=_ExperimentalConfig(verbose=True),
)
self.num_running_seqs = num_running_seqs
def __enter__(self):
return super().__enter__()
def __exit__(self, exc_type, exc_val, exc_tb):
super().__exit__(exc_type, exc_val, exc_tb)
self.results = LayerwiseProfileResults(
self.profiler.kineto_results, num_running_seqs=self.num_running_seqs
)
vllm/profiler/utils.py
-141
View File
@@ -1,141 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import dataclasses
from collections.abc import Callable
from _typeshed import DataclassInstance
from torch._C._profiler import _EventType, _ProfilerEvent, _TensorMetadata
#
# String / Print Manipulation
#
def trim_string_back(string: str, width: int) -> str:
if len(string) > width:
offset = len(string) - width + 3
string = string[:-offset]
if len(string) > 3:
string = string + "..."
return string
class TablePrinter:
def __init__(self, row_cls: type[DataclassInstance], column_widths: dict[str, int]):
self.row_cls = row_cls
self.fieldnames = [x.name for x in dataclasses.fields(row_cls)]
self.column_widths = column_widths
assert set(self.column_widths.keys()) == set(self.fieldnames)
def print_table(self, rows: list[DataclassInstance]):
self._print_header()
self._print_line()
for row in rows:
self._print_row(row)
def _print_header(self):
for i, f in enumerate(self.fieldnames):
last = i == len(self.fieldnames) - 1
col_width = self.column_widths[f]
print(
trim_string_back(f, col_width).ljust(col_width),
end=" | " if not last else "\n",
)
def _print_row(self, row):
assert isinstance(row, self.row_cls)
for i, f in enumerate(self.fieldnames):
last = i == len(self.fieldnames) - 1
col_width = self.column_widths[f]
val = getattr(row, f)
val_str = ""
if isinstance(val, str):
val_str = trim_string_back(val, col_width).ljust(col_width)
elif type(val) in [float, int]:
val_str = f"{float(val):>.2f}".rjust(col_width)
else:
val_str = f"{val}".rjust(col_width)
print(val_str, end=" | " if not last else "\n")
def _print_line(self):
total_col_width = 0
for column_width in self.column_widths.values():
total_col_width += column_width
print("=" * (total_col_width + 3 * (len(self.column_widths) - 1)))
def indent_string(
string: str, indent: int, indent_style: Callable[[int], str] | str = " "
) -> str:
if indent:
if isinstance(indent_style, str):
return indent_style * indent + string
else:
return indent_style(indent) + string
else:
return string
#
# _ProfilerEvent utils
#
def event_has_module(event: _ProfilerEvent) -> bool:
event_type, typed_event = event.typed
if event_type == _EventType.PyCall:
return typed_event.module is not None
return False
def event_is_torch_op(event: _ProfilerEvent) -> bool:
return event.tag == _EventType.TorchOp
def event_arg_repr(arg) -> str:
if arg is None or type(arg) in [float, int, bool, str]:
return f"{arg}"
elif isinstance(arg, list):
return f"[{', '.join([event_arg_repr(x) for x in arg])}]"
elif isinstance(arg, tuple):
return f"({', '.join([event_arg_repr(x) for x in arg])})"
else:
assert isinstance(arg, _TensorMetadata), f"Unsupported type: {type(arg)}"
sizes_str = ", ".join([str(x) for x in arg.sizes])
return f"{str(arg.dtype).replace('torch.', '')}[{sizes_str}]"
def event_torch_op_repr(event: _ProfilerEvent) -> str:
assert event.tag == _EventType.TorchOp
args_str = ", ".join([event_arg_repr(x) for x in event.typed[1].inputs])
return f"{event.name}({args_str})".replace("aten::", "")
def event_module_repr(event: _ProfilerEvent) -> str:
assert event_has_module(event)
module = event.typed[1].module
if module.parameters and len(module.parameters) > 0:
args_str = ", ".join(
[f"{x[0]}={event_arg_repr(x[1])}" for x in module.parameters]
)
return f"{module.cls_name}({args_str})"
else:
return module.cls_name
def event_torch_op_stack_trace(
curr_event: _ProfilerEvent, until: Callable[[_ProfilerEvent], bool]
) -> str:
trace = ""
curr_event = curr_event.parent
while curr_event and not until(curr_event):
if event_is_torch_op(curr_event):
if len(trace) > 0:
trace += " <- "
trace += event_torch_op_repr(curr_event)
curr_event = curr_event.parent
return trace
vllm/renderers/inkling_encoding.py
-383
View File
@@ -1,383 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Inkling chat-encoding core.
Pure implementation of Inkling chat rendering, kept deliberately free of
vLLM imports: it depends only on the :class:`InklingTextTokenizer` protocol
and speaks OpenAI-style message dicts. If a standalone Inkling
input-processing library becomes available, this module is the unit to
swap out — the swap point is marked in ``vllm/renderers/inkling.py``.
Multimodal parts follow the contract of vLLM's Inkling multimodal processor
(``InklingMultiModalProcessor`` anchors prompt updates on the bare
content-kind marker and inserts the per-patch placeholder run itself):
* image parts emit only ``<|content_image|>`` — no seed placeholder id;
* audio parts emit ``<|content_audio_input|><|audio_end|>`` — no seed
placeholder id between them.
``reasoning_effort`` (a float in [0, 0.99], sourced from
``chat_template_kwargs`` only) renders the ``Thinking effort level:``
system control block after the tool declarations and initial system messages.
"""
from __future__ import annotations
import json
from collections.abc import Iterator, Mapping, Sequence
from typing import Any, Protocol
END_OF_TEXT = "<|endoftext|>"
MESSAGE_USER = "<|message_user|>"
MESSAGE_MODEL = "<|message_model|>"
MESSAGE_SYSTEM = "<|message_system|>"
MESSAGE_TOOL = "<|message_tool|>"
CONTENT_TEXT = "<|content_text|>"
CONTENT_IMAGE = "<|content_image|>"
CONTENT_MODEL_END_SAMPLING = "<|content_model_end_sampling|>"
CONTENT_THINKING = "<|content_thinking|>"
CONTENT_AUDIO_INPUT = "<|content_audio_input|>"
CONTENT_TOOL_ERROR = "<|content_tool_error|>"
CONTENT_XML = "<|content_xml|>"
CONTENT_INVOKE_TOOL_JSON = "<|content_invoke_tool_json|>"
CONTENT_INVOKE_TOOL_TEXT = "<|content_invoke_tool_text|>"
END_MESSAGE = "<|end_message|>"
AUDIO_END = "<|audio_end|>"
ROLE_MESSAGE_TOKENS: dict[str, str] = {
"user": MESSAGE_USER,
"assistant": MESSAGE_MODEL,
"system": MESSAGE_SYSTEM,
# The developer role folds into the system role on the Inkling wire.
"developer": MESSAGE_SYSTEM,
"tool": MESSAGE_TOOL,
}
# Alternate vocab spellings per semantic token. The Inkling HF tokenizer
# exposes some semantic slots as ``<|unused_NNNNNN|>`` tokens (notably
# CONTENT_XML = <|unused_200024|>), so ID resolution must try these
# spellings in order.
SPECIAL_TOKEN_SPELLINGS: dict[str, tuple[str, ...]] = {
MESSAGE_USER: (MESSAGE_USER,),
MESSAGE_MODEL: (MESSAGE_MODEL,),
MESSAGE_SYSTEM: (MESSAGE_SYSTEM,),
MESSAGE_TOOL: (MESSAGE_TOOL,),
CONTENT_TEXT: (CONTENT_TEXT,),
CONTENT_IMAGE: (CONTENT_IMAGE,),
CONTENT_MODEL_END_SAMPLING: (CONTENT_MODEL_END_SAMPLING,),
CONTENT_THINKING: (CONTENT_THINKING,),
CONTENT_AUDIO_INPUT: (CONTENT_AUDIO_INPUT,),
CONTENT_XML: (CONTENT_XML, "<|unused_200024|>"),
CONTENT_INVOKE_TOOL_JSON: (CONTENT_INVOKE_TOOL_JSON,),
END_MESSAGE: (END_MESSAGE,),
AUDIO_END: (AUDIO_END,),
}
class InklingTextTokenizer(Protocol):
"""Structural tokenizer contract required by the renderer."""
def encode_text(self, text: str) -> list[int]: ...
def encode_special(self, token: str) -> int: ...
# OpenAI content-part type spellings that mean image / audio (rendering only
# needs the kind, not the bytes — the bytes are handled by the MM processor).
_IMAGE_PART_TYPES = frozenset({"image", "input_image", "image_url"})
_AUDIO_PART_TYPES = frozenset({"audio", "input_audio", "audio_url"})
_MAX_REASONING_EFFORT = 0.99
def render_inkling_messages(
messages: Sequence[Mapping[str, Any]],
tokenizer: InklingTextTokenizer,
*,
add_generation_prompt: bool = True,
tools: Sequence[Mapping[str, Any]] | None = None,
reasoning_effort: float | None = None,
) -> list[int]:
"""Render chat messages to Inkling input ids.
PURE renderer: emits Inkling framing plus bare media markers; media
encoding and placeholder expansion happen later in the MM processor.
``add_generation_prompt`` appends the assistant turn opener so the
model continues into the response.
"""
input_ids: list[int] = []
tool_call_id_to_name: dict[str, str] = {}
# Request-level tools plus per-developer-message tools (Rust renderer
# semantics) are declared in a single leading system block.
all_tools = list(tools or [])
for message in messages:
if message.get("role") == "developer":
all_tools.extend(message.get("tools") or [])
if all_tools:
_append_message(
input_ids,
tokenizer,
"system",
"xml",
_tool_declare_json(all_tools),
author_name="tool_declare",
)
for message in messages:
role = _expect_role(message)
if role not in {"system", "developer"} and reasoning_effort is not None:
_append_reasoning_effort(input_ids, tokenizer, reasoning_effort)
reasoning_effort = None
if role == "tool":
tool_name = message.get("name") or tool_call_id_to_name.get(
str(message.get("tool_call_id") or ""), ""
)
_append_message(
input_ids,
tokenizer,
"tool",
"text",
_flatten_text_content(message.get("content")),
author_name=str(tool_name),
)
continue
if role == "assistant":
reasoning_content = message.get("reasoning")
if reasoning_content is None:
reasoning_content = message.get("reasoning_content")
if reasoning_content:
if not isinstance(reasoning_content, str):
raise TypeError(
"assistant reasoning_content must be a string for "
"Inkling rendering"
)
_append_message(
input_ids,
tokenizer,
"assistant",
"thinking",
reasoning_content,
)
for kind, text in _iter_render_parts(message.get("content")):
_append_message(input_ids, tokenizer, role, kind, text)
if role == "assistant":
for tool_call in message.get("tool_calls") or []:
name, args = _tool_call_name_and_args(tool_call)
tool_call_id = _as_mapping(tool_call).get("id")
if tool_call_id:
tool_call_id_to_name[str(tool_call_id)] = name
_append_message(
input_ids,
tokenizer,
"assistant",
"invoke_tool_json",
_tool_call_json(name, args),
author_name=name,
)
input_ids.append(tokenizer.encode_special(CONTENT_MODEL_END_SAMPLING))
if reasoning_effort is not None:
_append_reasoning_effort(input_ids, tokenizer, reasoning_effort)
if add_generation_prompt:
input_ids.append(tokenizer.encode_special(MESSAGE_MODEL))
return input_ids
def _append_reasoning_effort(
input_ids: list[int],
tokenizer: InklingTextTokenizer,
reasoning_effort: float,
) -> None:
_append_message(
input_ids,
tokenizer,
"system",
"text",
_thinking_effort_text(reasoning_effort),
)
def _thinking_effort_text(reasoning_effort: Any) -> str:
if isinstance(reasoning_effort, bool) or not isinstance(
reasoning_effort, (int, float)
):
raise TypeError(
"Inkling reasoning_effort must be a number in [0.0, 0.99], "
f"got {type(reasoning_effort).__name__}"
)
value = float(reasoning_effort)
if not 0.0 <= value <= _MAX_REASONING_EFFORT:
raise ValueError(
f"Inkling reasoning_effort must be in [0.0, 0.99], got {value}"
)
effort_text = f"{value:.2f}".rstrip("0").rstrip(".")
if effort_text in {"0", "-0"}:
effort_text = "0.0"
return f"Thinking effort level: {effort_text}"
def _append_message(
input_ids: list[int],
tokenizer: InklingTextTokenizer,
role: str,
kind: str,
text: str,
*,
author_name: str | None = None,
) -> None:
input_ids.append(tokenizer.encode_special(ROLE_MESSAGE_TOKENS[role]))
if author_name:
input_ids.extend(tokenizer.encode_text(author_name))
if kind == "text":
input_ids.append(tokenizer.encode_special(CONTENT_TEXT))
input_ids.extend(tokenizer.encode_text(text))
elif kind == "image":
# Bare marker only: the MM processor replaces the marker with
# `marker + placeholder * num_patches` (see mm_preprocess.py).
input_ids.append(tokenizer.encode_special(CONTENT_IMAGE))
elif kind == "audio":
input_ids.append(tokenizer.encode_special(CONTENT_AUDIO_INPUT))
input_ids.append(tokenizer.encode_special(AUDIO_END))
elif kind == "thinking":
input_ids.append(tokenizer.encode_special(CONTENT_THINKING))
input_ids.extend(tokenizer.encode_text(text))
elif kind == "xml":
input_ids.append(tokenizer.encode_special(CONTENT_XML))
input_ids.extend(tokenizer.encode_text(text))
elif kind == "invoke_tool_json":
input_ids.append(tokenizer.encode_special(CONTENT_INVOKE_TOOL_JSON))
input_ids.extend(tokenizer.encode_text(text))
else:
raise ValueError(f"unsupported Inkling render part kind: {kind!r}")
input_ids.append(tokenizer.encode_special(END_MESSAGE))
def _iter_render_parts(content: Any) -> Iterator[tuple[str, str]]:
"""Yield (kind, text) per content part: kind in {text, image, audio}."""
if content is None:
return
if isinstance(content, str):
if content:
yield ("text", content)
return
if not isinstance(content, Sequence) or isinstance(content, (bytes, bytearray)):
raise TypeError("message content must be a string or a sequence of parts")
for part in content:
if isinstance(part, str):
yield ("text", part)
continue
if not isinstance(part, Mapping):
raise TypeError(f"content part must be mapping, got {type(part).__name__}")
ptype = part.get("type")
if ptype in (None, "text", "input_text"):
text = part.get("text", "")
yield ("text", text if isinstance(text, str) else "")
elif ptype in _IMAGE_PART_TYPES:
yield ("image", "")
elif ptype in _AUDIO_PART_TYPES:
yield ("audio", "")
else:
raise ValueError(f"unsupported content part type: {ptype!r}")
def _flatten_text_content(content: Any) -> str:
"""Flatten a tool-response content (string or text parts) to text."""
if content is None:
return ""
if isinstance(content, str):
return content
parts: list[str] = []
for kind, text in _iter_render_parts(content):
if kind != "text":
raise ValueError(
"Inkling tool response content must be text, "
f"got a part of kind {kind!r}"
)
parts.append(text)
return "".join(parts)
def _expect_role(message: Mapping[str, Any]) -> str:
role = message.get("role")
if role not in ROLE_MESSAGE_TOKENS:
raise ValueError(
f"unsupported Inkling message role {role!r}; "
f"expected one of {sorted(ROLE_MESSAGE_TOKENS)}"
)
return str(role)
def _as_mapping(value: Any) -> Mapping[str, Any]:
if isinstance(value, Mapping):
return value
if hasattr(value, "model_dump"):
dumped = value.model_dump()
if isinstance(dumped, Mapping):
return dumped
raise TypeError(f"expected mapping, got {type(value).__name__}")
def _canonical_json(value: Any) -> str:
return json.dumps(
_sort_json(value),
ensure_ascii=False,
allow_nan=False,
separators=(",", ":"),
)
def _sort_json(value: Any) -> Any:
if isinstance(value, Mapping):
return {str(key): _sort_json(value[key]) for key in sorted(value)}
if isinstance(value, Sequence) and not isinstance(value, (str, bytes, bytearray)):
return [_sort_json(item) for item in value]
return value
def _tool_declare_json(tools: Sequence[Mapping[str, Any]]) -> str:
tool_specs = []
for tool_value in tools:
tool = _as_mapping(tool_value)
function = _as_mapping(tool.get("function", {}))
tool_specs.append(
{
"description": function.get("description") or "",
"name": function["name"],
"parameters": function.get("parameters") or {},
"type": tool.get("type", "function"),
}
)
return _canonical_json(tool_specs)
def _tool_call_name_and_args(tool_call_value: Any) -> tuple[str, Mapping[str, Any]]:
tool_call = _as_mapping(tool_call_value)
function = _as_mapping(tool_call.get("function", {}))
name = function.get("name")
if not isinstance(name, str):
raise TypeError("tool call function name must be a string")
raw_args = function.get("arguments") or {}
if isinstance(raw_args, str):
args = json.loads(raw_args) if raw_args.strip() else {}
else:
args = raw_args
if not isinstance(args, Mapping):
raise TypeError("tool call function arguments must decode to an object")
return name, args
def _tool_call_json(name: str, args: Mapping[str, Any]) -> str:
name_json = json.dumps(name, ensure_ascii=False, allow_nan=False)
return f'{{"name":{name_json},"args":{_canonical_json(args)}}}'
vllm/snapshot/__init__.py
-4
View File
@@ -1,4 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Create and restore explicitly compatible vLLM engine snapshots."""
vllm/snapshot/controller.py
-132
View File
@@ -1,132 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Create and restore transaction orchestration for local snapshots."""
import argparse
import shutil
import sys
from pathlib import Path
from vllm.snapshot.manifest import (
read_manifest,
validate_artifact_root,
validate_identity,
)
from vllm.snapshot.runtime import (
LocalSnapshotTools,
SnapshotCreateError,
SnapshotRestoreError,
_error_detail,
)
from vllm.snapshot.types import Oracle, oracles_match
def _child_engine_argv(
args: argparse.Namespace, engine_argv: tuple[str, ...] | None
) -> tuple[str, ...]:
"""Build the child argv from the first argv source that is present.
In order: an explicit ``engine_argv`` from the caller, the process argv
tail when invoked as ``vllm snapshot create ...``, then the bare model from
the parsed args. Whichever source wins, snapshot options are stripped and
sleep mode is appended.
"""
if not engine_argv:
process_argv = tuple(sys.argv[1:])
if process_argv[:2] == ("snapshot", "create"):
engine_argv = process_argv[2:]
else:
model = getattr(args, "model_tag", None) or getattr(args, "model", None)
if not model:
raise SnapshotCreateError("snapshot create requires a model argument")
engine_argv = (str(model),)
child_argv: list[str] = []
iterator = iter(engine_argv)
for item in iterator:
option, separator, _ = item.partition("=")
option = option.replace("_", "-")
if len(option) > 2 and "--snapshot-dir".startswith(option):
if not separator:
next(iterator, None)
else:
child_argv.append(item)
if "--enable-sleep-mode" not in child_argv:
child_argv.append("--enable-sleep-mode")
return tuple(child_argv)
def create_snapshot(
args: argparse.Namespace,
*,
engine_argv: tuple[str, ...] | None = None,
tools: LocalSnapshotTools | None = None,
) -> None:
target = Path(args.snapshot_dir).absolute()
toolset = tools or LocalSnapshotTools()
toolset.preflight("create", target)
if target.exists() or target.is_symlink():
raise SnapshotCreateError(f"snapshot target already exists: {target}")
validate_artifact_root(target, creating=True)
target.parent.mkdir(mode=0o700, parents=False, exist_ok=True)
validate_artifact_root(target.parent, creating=False)
target.mkdir(mode=0o700)
validate_artifact_root(target, creating=False)
published = False
root_pid: int | None = None
try:
child_argv = _child_engine_argv(args, engine_argv)
root_pid = toolset.launch_child(target, child_argv)
oracle = toolset.wait_ready(target, root_pid)
inventory = toolset.inventory(root_pid, target)
toolset.dump(target, inventory)
toolset.verify_dead(inventory)
manifest = toolset.make_manifest(args, child_argv, inventory, oracle, target)
toolset.publish(target, manifest)
published = True
except BaseException:
if root_pid is not None:
toolset.abort_create(root_pid)
raise
finally:
# Runs even when abort_create raised, so no partial artifact is left.
if not published and target.exists():
shutil.rmtree(target)
def restore_snapshot(
args: argparse.Namespace, *, tools: LocalSnapshotTools | None = None
) -> None:
artifact = Path(args.snapshot_dir).absolute()
toolset = tools or LocalSnapshotTools()
toolset.preflight("restore", artifact)
manifest = read_manifest(artifact)
validate_identity(manifest, toolset.current_identity(manifest.gpu_uuid))
root_pid: int | None = None
try:
root_pid = toolset.restore(artifact, manifest)
toolset.release(artifact, args.host, args.port)
toolset.wait_listener(root_pid, args.host, args.port)
actual = toolset.request_oracle(args.host, args.port, manifest)
expected = Oracle(
token_ids=manifest.oracle_token_ids,
text=manifest.oracle_text,
sampled_token_logprob=manifest.oracle_sampled_token_logprob,
)
if not oracles_match(expected, actual):
raise SnapshotRestoreError(
f"snapshot oracle mismatch: expected={expected!r}, actual={actual!r}"
)
toolset.complete_restore(root_pid)
except BaseException as error:
if root_pid is not None:
try:
toolset.cleanup(root_pid)
except BaseException as cleanup_error:
raise SnapshotRestoreError(
f"snapshot restore failed: {_error_detail(error)}; "
f"cleanup failed: {_error_detail(cleanup_error)}"
) from error
if isinstance(error, SnapshotRestoreError):
raise
raise SnapshotRestoreError(str(error)) from error
vllm/snapshot/manifest.py
-276
View File
@@ -1,276 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import json
import math
import os
import stat
from pathlib import Path
from typing import Annotated, Any, Literal, Self
from pydantic import (
BaseModel,
ConfigDict,
Field,
StrictFloat,
StrictInt,
StrictStr,
ValidationError,
field_validator,
model_validator,
)
class SnapshotCompatibilityError(RuntimeError):
"""The snapshot identity does not match the requested runtime."""
class SnapshotSecurityError(RuntimeError):
"""The snapshot artifact does not meet the private-path contract."""
_MAX_PID = 2**31 - 1
_PositivePid = Annotated[int, Field(gt=0, le=_MAX_PID)]
_OracleTokenId = Annotated[int, Field(ge=0)]
class SnapshotRuntimeIdentity(BaseModel):
model_config = ConfigDict(extra="forbid", frozen=True, strict=True)
vllm_version: str
python_version: str
torch_version: str
cuda_runtime: str
driver_version: str
criu_version: str
cuda_checkpoint_sha256: str
kernel_release: str
host_id: str
gpu_name: str
gpu_uuid: str
environment: tuple[tuple[str, str], ...]
class SnapshotManifest(SnapshotRuntimeIdentity):
schema_version: Literal[1]
boundary: Literal["post-engine-init-reloadable-state-released"]
created_at: str
artifact_bytes: Annotated[int, Field(ge=0)]
model: str
served_model_name: str
model_revision: str
tokenizer_revision: str
engine_argv: tuple[str, ...]
process_tree: Annotated[tuple[_PositivePid, ...], Field(min_length=1)]
cuda_holders: Annotated[tuple[_PositivePid, ...], Field(min_length=1)]
external_cache_files: tuple[tuple[str, str], ...] = ()
oracle_token_ids: Annotated[
tuple[_OracleTokenId, ...], Field(min_length=1, max_length=1)
]
oracle_text: str
oracle_sampled_token_logprob: float
@field_validator("schema_version", mode="before")
@classmethod
def _require_schema_version_one(cls, value: Any) -> Any:
if type(value) is not int or value != 1:
raise ValueError("schema version must be integer 1")
return value
@field_validator("oracle_sampled_token_logprob", mode="before")
@classmethod
def _require_finite_float(cls, value: Any) -> Any:
if type(value) is not float or not math.isfinite(value):
raise ValueError("sampled token log probability must be a finite float")
return value
@model_validator(mode="after")
def _validate_process_relationships(self) -> Self:
if len(set(self.process_tree)) != len(self.process_tree):
raise ValueError("process_tree contains duplicate PIDs")
if len(set(self.cuda_holders)) != len(self.cuda_holders):
raise ValueError("cuda_holders contains duplicate PIDs")
if not set(self.cuda_holders).issubset(self.process_tree):
raise ValueError("cuda_holders must be a subset of process_tree")
return self
class ReleaseMarker(BaseModel):
"""Release marker the controller writes; unknown keys are ignored."""
release: Literal[True]
host: StrictStr | None = None
port: Annotated[StrictInt, Field(ge=1, le=65535)]
@field_validator("release", mode="before")
@classmethod
def _require_json_true(cls, value: Any) -> Any:
# A literal accepts 1 as True; the barrier requires JSON true itself.
if value is not True:
raise ValueError("release must be JSON true")
return value
class ReadyMarker(BaseModel):
"""Ready marker the child writes; unknown keys are ignored."""
token_ids: tuple[StrictInt, ...]
text: StrictStr
sampled_token_logprob: StrictFloat | StrictInt
def validate_identity(
expected: SnapshotRuntimeIdentity, actual: SnapshotRuntimeIdentity
) -> None:
"""Require an exact match for every field that can affect compatibility."""
# ``expected`` may be a manifest, so compare the identity fields only.
fields = set(SnapshotRuntimeIdentity.model_fields)
expected_values = expected.model_dump(include=fields)
actual_values = actual.model_dump(include=fields)
differing = [
name
for name in SnapshotRuntimeIdentity.model_fields
if expected_values[name] != actual_values[name]
]
if differing:
raise SnapshotCompatibilityError(f"snapshot mismatch: {', '.join(differing)}")
def _existing_path_chain(path: Path):
current = path
while True:
if current.exists() or current.is_symlink():
yield current
if current.parent == current:
break
current = current.parent
def _is_trusted_sticky_directory(path_stat: os.stat_result) -> bool:
return (
path_stat.st_uid == 0
and stat.S_ISDIR(path_stat.st_mode)
and bool(path_stat.st_mode & stat.S_ISVTX)
)
def validate_artifact_root(path: Path, *, creating: bool) -> None:
"""Validate that an artifact path cannot be replaced by another user."""
path = Path(os.path.abspath(path))
if path.is_symlink():
raise SnapshotSecurityError(f"snapshot path is a symlink: {path}")
if not creating and not path.exists():
raise SnapshotSecurityError(f"snapshot path does not exist: {path}")
effective_uid = os.geteuid()
for component in _existing_path_chain(path):
component_stat = component.lstat()
if stat.S_ISLNK(component_stat.st_mode):
raise SnapshotSecurityError(
f"snapshot path contains a symlink: {component}"
)
if component_stat.st_uid not in (0, effective_uid):
raise SnapshotSecurityError(
f"snapshot path has an untrusted owner: {component}"
)
if component_stat.st_mode & (
stat.S_IWGRP | stat.S_IWOTH
) and not _is_trusted_sticky_directory(component_stat):
raise SnapshotSecurityError(
f"snapshot path has a group- or world-writable ancestor: {component}"
)
if path.exists():
path_stat = path.lstat()
if not stat.S_ISDIR(path_stat.st_mode):
raise SnapshotSecurityError(f"snapshot path is not a directory: {path}")
if path_stat.st_uid != effective_uid:
raise SnapshotSecurityError(
f"snapshot directory is not owned by the current user: {path}"
)
if stat.S_IMODE(path_stat.st_mode) & 0o077:
raise SnapshotSecurityError(
f"snapshot directory must have mode 0700: {path}"
)
def _validation_path(error: ValidationError) -> str:
detail = error.errors(include_url=False)[0]
location = detail["loc"]
if not location:
return "JSON" if detail["type"] == "json_invalid" else "root"
return ".".join(map(str, location))
def _fsync_directory(path: Path) -> None:
directory_fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
try:
os.fsync(directory_fd)
finally:
os.close(directory_fd)
def _write_json_atomic(path: Path, payload: object, *, overwrite: bool = False) -> None:
temporary = path.with_name(f"{path.name}.tmp")
if overwrite:
temporary.unlink(missing_ok=True)
elif path.exists() or path.is_symlink():
raise FileExistsError(path)
flags = os.O_WRONLY | os.O_CREAT | os.O_EXCL | getattr(os, "O_NOFOLLOW", 0)
descriptor = os.open(temporary, flags, 0o600)
try:
with os.fdopen(descriptor, "wb") as output:
output.write(
json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
)
output.flush()
os.fsync(output.fileno())
os.replace(temporary, path)
_fsync_directory(path.parent)
except BaseException:
temporary.unlink(missing_ok=True)
raise
def write_manifest_atomic(path: Path, manifest: SnapshotManifest) -> None:
"""Write a new private manifest and make its directory entry durable."""
path = Path(path)
validate_artifact_root(path, creating=False)
destination = path / "manifest.json"
try:
_write_json_atomic(destination, manifest.model_dump(mode="json"))
except FileExistsError as error:
raise SnapshotSecurityError(
f"manifest already exists: {destination}"
) from error
def read_manifest(path: Path) -> SnapshotManifest:
"""Read a snapshot published by atomically installing its manifest."""
path = Path(path)
validate_artifact_root(path, creating=False)
manifest_path = path / "manifest.json"
if manifest_path.is_symlink():
raise SnapshotSecurityError(f"manifest is a symlink: {manifest_path}")
flags = os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)
try:
file_descriptor = os.open(manifest_path, flags)
except FileNotFoundError as error:
raise SnapshotSecurityError(
f"manifest does not exist: {manifest_path}"
) from error
try:
with os.fdopen(file_descriptor, encoding="utf-8") as manifest_file:
payload = manifest_file.read()
except UnicodeDecodeError as error:
raise SnapshotCompatibilityError("invalid snapshot manifest: JSON") from error
try:
return SnapshotManifest.model_validate_json(payload, strict=True)
except ValidationError as error:
raise SnapshotCompatibilityError(
f"invalid snapshot manifest: {_validation_path(error)}"
) from error
def inspect_snapshot(path: Path) -> dict[str, object]:
return read_manifest(path).model_dump(mode="json")
vllm/snapshot/runtime.py
-1318
View File
File diff suppressed because it is too large Load Diff
vllm/snapshot/server.py
-252
View File
@@ -1,252 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import asyncio
import json
import os
import stat
import sys
from dataclasses import dataclass
from pathlib import Path
from typing import Any
from pydantic import ValidationError
from vllm.snapshot.manifest import ReleaseMarker, _validation_path, _write_json_atomic
from vllm.snapshot.types import Oracle, oracles_match
_CANARY_PROMPT = "The capital of France is"
class SnapshotCanaryError(RuntimeError):
"""The initialized engine did not produce a valid snapshot oracle."""
class SnapshotBarrierError(RuntimeError):
"""The controller supplied an invalid snapshot release marker."""
@dataclass(frozen=True)
class ListenerConfig:
host: str | None
port: int
@dataclass(frozen=True)
class ControlArgs:
ready_file: Path
release_file: Path
release_timeout_s: float
def write_ready_atomic(path: Path, oracle: Oracle) -> None:
_write_json_atomic(
path,
{
"sampled_token_logprob": oracle.sampled_token_logprob,
"token_ids": oracle.token_ids,
"text": oracle.text,
},
)
def read_release_marker(path: Path) -> ListenerConfig:
path = Path(path)
path_stat = path.lstat()
if stat.S_ISLNK(path_stat.st_mode):
raise SnapshotBarrierError("release marker must not be a symlink")
try:
payload = json.loads(path.read_text())
except (OSError, json.JSONDecodeError) as error:
raise SnapshotBarrierError("release marker is not valid JSON") from error
try:
marker = ReleaseMarker.model_validate(payload)
except ValidationError as error:
raise SnapshotBarrierError(
f"release marker is invalid: {_validation_path(error)}"
) from error
return ListenerConfig(host=marker.host, port=marker.port)
def parse_control_args(argv: list[str]) -> tuple[ControlArgs, list[str]]:
parser = argparse.ArgumentParser(add_help=False)
parser.add_argument("--ready-file", type=Path, required=True)
parser.add_argument("--release-file", type=Path, required=True)
parser.add_argument("--release-timeout-s", type=float, default=900.0)
control_args, remaining = parser.parse_known_args(argv)
if remaining and remaining[0] == "--":
remaining = remaining[1:]
if control_args.release_timeout_s <= 0:
raise ValueError("release timeout must be positive")
return (
ControlArgs(
ready_file=control_args.ready_file,
release_file=control_args.release_file,
release_timeout_s=control_args.release_timeout_s,
),
remaining,
)
async def _release_reloadable_state(engine: Any) -> None:
"""Discard model and KV state before the process image is captured."""
await engine.sleep(level=2)
async def _restore_reloadable_state(engine: Any) -> None:
"""Rebuild state discarded by ``_release_reloadable_state``."""
await engine.wake_up(tags=["weights"])
await engine.collective_rpc("reload_weights")
await engine.wake_up(tags=["kv_cache"])
def oracle_from_request_output(request_output: Any) -> Oracle:
try:
candidate = request_output.outputs[0]
token_ids = tuple(candidate.token_ids)
(sampled_token_id,) = token_ids
return Oracle(
token_ids=token_ids,
text=candidate.text,
sampled_token_logprob=candidate.logprobs[0][sampled_token_id].logprob,
)
except ValueError as error:
raise SnapshotCanaryError(
"snapshot canary must produce exactly one finite token logprob"
) from error
except (AttributeError, IndexError, KeyError, TypeError) as error:
raise SnapshotCanaryError(
"snapshot canary did not return sampled token logprob"
) from error
async def run_engine_canary(engine: Any) -> Oracle:
from vllm import SamplingParams
final_output = None
sampling_params = SamplingParams(
temperature=0,
min_tokens=1,
max_tokens=1,
seed=0,
logprobs=0,
)
async for output in engine.generate(
_CANARY_PROMPT,
sampling_params,
request_id="vllm-snapshot-canary",
):
final_output = output
if final_output is None:
raise SnapshotCanaryError("canary generation returned no output")
return oracle_from_request_output(final_output)
def detach_snapshot_streams() -> None:
"""Remove launch-log descriptors from the reusable process image."""
sys.stdout.flush()
sys.stderr.flush()
sink = os.open(os.devnull, os.O_WRONLY)
try:
os.dup2(sink, sys.stdout.fileno())
os.dup2(sink, sys.stderr.fileno())
finally:
os.close(sink)
async def wait_for_release_marker(
path: Path,
*,
timeout_s: float,
poll_interval_s: float = 0.05,
) -> ListenerConfig:
if timeout_s <= 0:
raise ValueError("release timeout must be positive")
if poll_interval_s <= 0:
raise ValueError("release poll interval must be positive")
remaining_s = timeout_s
while not path.exists():
if remaining_s <= 0:
raise TimeoutError(f"release marker not found before timeout: {path}")
delay_s = min(poll_interval_s, remaining_s)
await asyncio.sleep(delay_s)
remaining_s -= delay_s
return read_release_marker(path)
def parse_vllm_args(argv: list[str]) -> Any:
from vllm.entrypoints.launchers.cli_args import (
make_arg_parser,
validate_parsed_serve_args,
)
from vllm.utils.argparse_utils import FlexibleArgumentParser
parser = FlexibleArgumentParser(prog="python -m vllm.snapshot.server")
args = make_arg_parser(parser).parse_args(argv)
if args.model_tag is not None:
args.model = args.model_tag
if args.grpc or args.headless:
raise ValueError("snapshot server requires the HTTP frontend")
if args.api_server_count not in (None, 1):
raise ValueError("snapshot server supports one API server")
args.api_server_count = None
validate_parsed_serve_args(args)
return args
async def run_vllm_snapshot_child(control: ControlArgs, args: Any) -> None:
from vllm.entrypoints.launchers.api_server.entry import (
build_and_serve,
build_async_engine_client,
)
from vllm.entrypoints.launchers.launcher import setup_server
args.enable_sleep_mode = True
async with build_async_engine_client(args) as engine:
oracle = await run_engine_canary(engine)
try:
await _release_reloadable_state(engine)
await _restore_reloadable_state(engine)
rehearsal_oracle = await run_engine_canary(engine)
if not oracles_match(oracle, rehearsal_oracle):
raise SnapshotCanaryError("snapshot rehearsal changed canary output")
await _release_reloadable_state(engine)
except SnapshotCanaryError:
raise
except Exception as error:
raise SnapshotCanaryError("snapshot rehearsal failed") from error
detach_snapshot_streams()
write_ready_atomic(control.ready_file, oracle)
listener = await wait_for_release_marker(
control.release_file,
timeout_s=control.release_timeout_s,
)
args.host = listener.host
args.port = listener.port
await _restore_reloadable_state(engine)
listen_address, sock = setup_server(args, reuse_port=False)
try:
shutdown_task = await build_and_serve(
engine,
listen_address,
sock,
args,
)
await shutdown_task
finally:
sock.close()
def main(argv: list[str] | None = None) -> None:
control, remaining = parse_control_args(sys.argv[1:] if argv is None else argv)
args = parse_vllm_args(remaining)
import uvloop
uvloop.run(run_vllm_snapshot_child(control, args))
if __name__ == "__main__":
main()
vllm/snapshot/types.py
-37
View File
@@ -1,37 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import math
from dataclasses import dataclass
@dataclass(frozen=True)
class Oracle:
token_ids: tuple[int, ...]
text: str
sampled_token_logprob: float
def __post_init__(self) -> None:
if len(self.token_ids) != 1:
raise ValueError("snapshot oracle must contain exactly one token")
value = self.sampled_token_logprob
if (
isinstance(value, bool)
or not isinstance(value, (int, float))
or not math.isfinite(value)
):
raise ValueError("sampled token logprob must be finite")
object.__setattr__(self, "sampled_token_logprob", float(value))
def oracles_match(expected: Oracle, actual: Oracle) -> bool:
return (
expected.token_ids == actual.token_ids
and expected.text == actual.text
and math.isclose(
expected.sampled_token_logprob,
actual.sampled_token_logprob,
rel_tol=0.0,
abs_tol=1e-3,
)
)
vllm/transformers_utils/processors/deepseek_ocr.py
-459
View File
@@ -1,459 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# adapted from https://github.com/deepseek-ai/DeepSeek-OCR/blob/main/DeepSeek-OCR-master/DeepSeek-OCR-vllm/process/image_process.py
# and https://github.com/deepseek-ai/DeepSeek-OCR-2/blob/main/DeepSeek-OCR2-master/DeepSeek-OCR2-vllm/process/image_process.py
import math
from typing import Literal
import torch
import torchvision.transforms as T
from PIL import Image, ImageOps
from transformers import BatchFeature, LlamaTokenizerFast
from transformers.processing_utils import ProcessorMixin
# TODO(Isotr0py): change modes for variants
# see: https://github.com/deepseek-ai/DeepSeek-OCR/blob/8cf003d38821fa1b19c73da3bd1b0dc262ea8136/DeepSeek-OCR-master/DeepSeek-OCR-vllm/config.py#L1-L6
# Tiny: base_size = 512, image_size = 512, crop_mode = False
# Small: base_size = 640, image_size = 640, crop_mode = False
# Base: base_size = 1024, image_size = 1024, crop_mode = False
# Large: base_size = 1280, image_size = 1280, crop_mode = False
# Gundam: base_size = 1024, image_size = 640, crop_mode = True
BASE_SIZE = 1024
IMAGE_SIZE = 640
CROP_MODE = True
# TODO(Isotr0py): Expose as mm_kwargs
MIN_CROPS = 2
MAX_CROPS = 6 # max:9; If your GPU memory is small, it is recommended to set it to 6.
def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
best_ratio_diff = float("inf")
best_ratio = (1, 1)
area = width * height
for ratio in target_ratios:
target_aspect_ratio = ratio[0] / ratio[1]
ratio_diff = abs(aspect_ratio - target_aspect_ratio)
if ratio_diff < best_ratio_diff:
best_ratio_diff = ratio_diff
best_ratio = ratio
elif ratio_diff == best_ratio_diff:
if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
best_ratio = ratio
return best_ratio
def calculate_aspect_ratios(
min_num: int = MIN_CROPS, max_num: int = MAX_CROPS
) -> list[tuple[int, int]]:
target_ratios: set[tuple[int, int]] = set(
(i, j)
for n in range(min_num, max_num + 1)
for i in range(1, n + 1)
for j in range(1, n + 1)
if i * j <= max_num and i * j >= min_num
)
sorted_target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
return sorted_target_ratios
def count_tiles(
orig_width,
orig_height,
min_num=MIN_CROPS,
max_num=MAX_CROPS,
image_size=640,
use_thumbnail=False,
):
aspect_ratio = orig_width / orig_height
# calculate the existing image aspect ratio
target_ratios = calculate_aspect_ratios(min_num, max_num)
# find the closest aspect ratio to the target
target_aspect_ratio = find_closest_aspect_ratio(
aspect_ratio, target_ratios, orig_width, orig_height, image_size
)
return target_aspect_ratio
def dynamic_preprocess(
image, min_num=MIN_CROPS, max_num=MAX_CROPS, image_size=640, use_thumbnail=False
):
orig_width, orig_height = image.size
aspect_ratio = orig_width / orig_height
# calculate the existing image aspect ratio
target_ratios = calculate_aspect_ratios(min_num, max_num)
# find the closest aspect ratio to the target
target_aspect_ratio = find_closest_aspect_ratio(
aspect_ratio, target_ratios, orig_width, orig_height, image_size
)
# calculate the target width and height
target_width = image_size * target_aspect_ratio[0]
target_height = image_size * target_aspect_ratio[1]
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
# resize the image
resized_img = image.resize((target_width, target_height))
processed_images = []
for i in range(blocks):
box = (
(i % (target_width // image_size)) * image_size,
(i // (target_width // image_size)) * image_size,
((i % (target_width // image_size)) + 1) * image_size,
((i // (target_width // image_size)) + 1) * image_size,
)
# split the image
split_img = resized_img.crop(box)
processed_images.append(split_img)
assert len(processed_images) == blocks
if use_thumbnail and len(processed_images) != 1:
thumbnail_img = image.resize((image_size, image_size))
processed_images.append(thumbnail_img)
return processed_images, target_aspect_ratio
class ImageTransform:
def __init__(
self,
mean: tuple[float, float, float] = (0.5, 0.5, 0.5),
std: tuple[float, float, float] = (0.5, 0.5, 0.5),
normalize: bool = True,
):
self.mean = mean
self.std = std
self.normalize = normalize
transform_pipelines = [T.ToTensor()]
if normalize:
transform_pipelines.append(T.Normalize(mean, std))
self.transform = T.Compose(transform_pipelines)
def __call__(self, pil_img: Image.Image):
x = self.transform(pil_img)
return x
class DeepseekOCRProcessor(ProcessorMixin):
tokenizer_class = ("LlamaTokenizer", "LlamaTokenizerFast")
attributes = ["tokenizer"]
def __init__(
self,
tokenizer: LlamaTokenizerFast,
patch_size: int = 16,
downsample_ratio: int = 4,
image_mean: tuple[float, float, float] = (0.5, 0.5, 0.5),
image_std: tuple[float, float, float] = (0.5, 0.5, 0.5),
normalize: bool = True,
image_token: str = "<image>",
pad_token: str = "<｜▁pad▁｜>",
add_special_token: bool = False,
sft_format: str = "deepseek",
mask_prompt: bool = True,
ignore_id: int = -100,
image_size: int = IMAGE_SIZE,
base_size: int = BASE_SIZE,
strategy: Literal["v1", "v2"] = "v1",
max_crops: int = MAX_CROPS,
**kwargs,
):
self.image_size = image_size
self.base_size = base_size
self.max_crops = max_crops
# image token calculation strategy for
# Deepseek-OCR and Deepseek-OCR-2
self.strategy = strategy
assert strategy in ["v1", "v2"], "Only 'v1' and 'v2' strategies are supported."
self.patch_size = 16
self.image_mean = image_mean
self.image_std = image_std
self.normalize = normalize
self.downsample_ratio = 4
self.image_transform = ImageTransform(
mean=image_mean, std=image_std, normalize=normalize
)
self.tokenizer = tokenizer
self.tokenizer.padding_side = "left" # must set this，padding side with make a difference in batch inference # noqa: E501
# add the pad_token as special token to use 'tokenizer.pad_token'
# and 'tokenizer.pad_token_id'
if self.tokenizer.pad_token is None:
self.tokenizer.add_special_tokens({"pad_token": pad_token})
# add image token
self.image_token_id = self.tokenizer.vocab.get(image_token)
self.image_token = image_token
self.pad_token = pad_token
self.add_special_token = add_special_token
self.sft_format = sft_format
self.mask_prompt = mask_prompt
self.ignore_id = ignore_id
super().__init__(
tokenizer,
**kwargs,
)
@property
def bos_id(self):
return self.tokenizer.bos_token_id
@property
def eos_id(self):
return self.tokenizer.eos_token_id
@property
def pad_id(self):
return self.tokenizer.pad_token_id
def encode(self, text: str, bos: bool = True, eos: bool = False):
t = self.tokenizer.encode(text, add_special_tokens=False)
if bos:
t = [self.bos_id] + t
if eos:
t = t + [self.eos_id]
return t
def decode(self, t: list[int], **kwargs) -> str:
return self.tokenizer.decode(t, **kwargs)
def process_one(
self,
prompt: str,
images: list[Image.Image],
crop_mode: bool = CROP_MODE,
):
"""Args:
prompt (str): the formatted prompt;
images (List[ImageType]): the list of images;
crop_mode (bool): if True, then crop the image;
Returns:
outputs (BaseProcessorOutput): the output of the processor,
- input_ids (torch.LongTensor): [N + image tokens]
- target_ids (torch.LongTensor): [N + image tokens]
- pixel_values (torch.FloatTensor): [n_patches, 3, H, W]
- image_id (int): the id of the image token
- num_image_tokens (List[int]): the number of image tokens
"""
if prompt is None or images is None:
raise ValueError("prompt and images must be used at the same time.")
sft_format = prompt
(
input_ids,
pixel_values,
images_crop,
images_seq_mask,
images_spatial_crop,
num_image_tokens,
_,
) = self.tokenize_with_images(
conversation=sft_format,
images=images,
bos=True,
eos=True,
cropping=crop_mode,
)
prepare = BatchFeature(
data=dict(
input_ids=input_ids,
pixel_values=pixel_values,
images_crop=images_crop,
images_seq_mask=images_seq_mask,
images_spatial_crop=images_spatial_crop,
num_image_tokens=num_image_tokens,
),
tensor_type="pt",
)
return prepare
def __call__(
self,
*,
prompt: str,
images: list[Image.Image],
crop_mode: bool = CROP_MODE,
**kwargs,
):
prepare = self.process_one(
prompt=prompt,
images=images,
crop_mode=crop_mode,
)
return prepare
def tokenize_with_images(
self,
conversation: str,
images: list[Image.Image],
bos: bool = True,
eos: bool = True,
cropping: bool = True,
):
"""Tokenize text with <image> tags."""
num_image_tags = conversation.count(self.image_token)
if num_image_tags != len(images):
raise ValueError(
f"Number of {self.image_token!r} tokens in prompt "
f"({num_image_tags}) does not match number of images "
f"({len(images)})."
)
text_splits = conversation.split(self.image_token)
images_list, images_crop_list, images_seq_mask, images_spatial_crop = (
[],
[],
[],
[],
)
image_shapes = []
num_image_tokens = []
tokenized_str = []
for text_sep, image in zip(text_splits, images):
tokenized_sep = self.encode(text_sep, bos=False, eos=False)
tokenized_str += tokenized_sep
images_seq_mask += [False] * len(tokenized_sep)
image_shapes.append(image.size)
images_crop_raw = []
if image.size[0] <= self.image_size and image.size[1] <= self.image_size:
crop_ratio = [1, 1]
elif cropping:
images_crop_raw, crop_ratio = dynamic_preprocess(
image, image_size=self.image_size, max_num=self.max_crops
)
else:
crop_ratio = [1, 1]
if not cropping:
image = image.resize((self.image_size, self.image_size))
global_view = ImageOps.pad(
image,
(self.base_size, self.base_size),
color=tuple(int(x * 255) for x in self.image_transform.mean),
)
images_list.append(self.image_transform(global_view))
num_width_tiles, num_height_tiles = crop_ratio
images_spatial_crop.append([num_width_tiles, num_height_tiles])
if num_width_tiles > 1 or num_height_tiles > 1:
for cropped_image in images_crop_raw:
images_crop_list.append(self.image_transform(cropped_image))
num_queries = math.ceil(
(self.image_size // self.patch_size) / self.downsample_ratio
)
num_queries_base = math.ceil(
(self.base_size // self.patch_size) / self.downsample_ratio
)
num_tokens_base = (
(num_queries_base * (num_queries_base + 1))
if self.strategy == "v1"
else num_queries_base * num_queries_base
)
tokenized_image = [self.image_token_id] * num_tokens_base
tokenized_image += [self.image_token_id]
if num_width_tiles > 1 or num_height_tiles > 1:
num_tokens_per_row = (
num_queries * num_width_tiles + 1
if self.strategy == "v1"
else num_queries * num_width_tiles
)
local_row = [self.image_token_id] * num_tokens_per_row
tokenized_image += local_row * (num_queries * num_height_tiles)
tokenized_str += tokenized_image
images_seq_mask += [True] * len(tokenized_image)
num_image_tokens.append(len(tokenized_image))
"""process the last text split"""
tokenized_sep = self.encode(text_splits[-1], bos=False, eos=False)
tokenized_str += tokenized_sep
images_seq_mask += [False] * len(tokenized_sep)
"""add the bos and eos tokens"""
if bos:
tokenized_str = [self.bos_id] + tokenized_str
images_seq_mask = [False] + images_seq_mask
if eos:
tokenized_str = tokenized_str + [self.eos_id]
images_seq_mask = images_seq_mask + [False]
assert len(tokenized_str) == len(images_seq_mask), (
f"tokenize_with_images func: tokenized_str's length {len(tokenized_str)} "
f"is not equal to images_seq_mask's length {len(images_seq_mask)}."
)
masked_tokenized_str = []
for token_index in tokenized_str:
if token_index != self.image_token_id:
masked_tokenized_str.append(token_index)
else:
masked_tokenized_str.append(self.ignore_id)
assert (
len(tokenized_str) == len(images_seq_mask) == len(masked_tokenized_str)
), (
f"tokenized_str's length {len(tokenized_str)}, "
f"input_ids' length {len(masked_tokenized_str)}, "
f"images_seq_mask's length {len(images_seq_mask)}, are not equal."
)
input_ids = torch.LongTensor(tokenized_str)
target_ids = torch.LongTensor(masked_tokenized_str)
images_seq_mask = torch.tensor(images_seq_mask, dtype=torch.bool)
# set input_ids < 0 | input_ids == self.image_token_id as ignore_id
target_ids[(input_ids < 0) | (input_ids == self.image_token_id)] = (
self.ignore_id
)
input_ids[input_ids < 0] = self.pad_id
# Remove the ending eos token
assert input_ids[-1] == self.eos_id
input_ids = input_ids[:-1]
target_ids = target_ids[:-1]
images_seq_mask = images_seq_mask[:-1]
if len(images_list) == 0:
pixel_values = torch.zeros((0, 3, self.base_size, self.base_size))
images_spatial_crop = torch.zeros((0, 2), dtype=torch.long)
images_crop = torch.zeros((0, 3, self.image_size, self.image_size))
else:
pixel_values = torch.stack(images_list, dim=0)
images_spatial_crop = torch.tensor(images_spatial_crop, dtype=torch.long)
if images_crop_list:
images_crop = torch.stack(images_crop_list, dim=0)
else:
images_crop = torch.zeros((0, 3, self.image_size, self.image_size))
input_ids = input_ids.unsqueeze(0)
return (
input_ids,
pixel_values,
images_crop,
images_seq_mask,
images_spatial_crop,
num_image_tokens,
image_shapes,
)
vllm/transformers_utils/processors/internvl.py
-585
View File
@@ -1,585 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# adapted from https://huggingface.co/OpenGVLab/InternVL2-4B/blob/main/modeling_internvl_chat.py
# --------------------------------------------------------
# InternVL
# Copyright (c) 2023 OpenGVLab
# Licensed under The MIT License [see LICENSE for details]
# --------------------------------------------------------
import numpy.typing as npt
import torch
import torchvision.transforms as T
from PIL import Image
from transformers import (
BaseVideoProcessor,
BatchFeature,
ImageProcessingMixin,
TensorType,
)
from transformers.processing_utils import ProcessorMixin
from vllm.multimodal.image import convert_image_mode
from vllm.multimodal.processing import PromptUpdateDetails, cached_encode
from vllm.tokenizers.hf import HfTokenizer
IMAGENET_MEAN = (0.485, 0.456, 0.406)
IMAGENET_STD = (0.229, 0.224, 0.225)
# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B
def build_transform(input_size: int):
MEAN, STD = IMAGENET_MEAN, IMAGENET_STD
return T.Compose(
[
T.Lambda(lambda img: convert_image_mode(img, "RGB")),
T.Resize(
(input_size, input_size), interpolation=T.InterpolationMode.BICUBIC
),
T.ToTensor(),
T.Normalize(mean=MEAN, std=STD),
]
)
# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B
def find_closest_aspect_ratio(
aspect_ratio: float,
target_ratios: list[tuple[int, int]],
*,
width: int,
height: int,
image_size: int,
) -> tuple[int, int]:
best_ratio_diff = float("inf")
best_ratio = (1, 1)
area = width * height
for ratio in target_ratios:
target_aspect_ratio = ratio[0] / ratio[1]
ratio_diff = abs(aspect_ratio - target_aspect_ratio)
if ratio_diff < best_ratio_diff:
best_ratio_diff = ratio_diff
best_ratio = ratio
elif ratio_diff == best_ratio_diff:
if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
best_ratio = ratio
return best_ratio
def resolve_internvl_min_max_num(
*,
min_dynamic_patch: int,
max_dynamic_patch: int,
dynamic_image_size: bool,
use_thumbnail: bool,
) -> tuple[int, int]:
min_dynamic_patch = min_dynamic_patch if dynamic_image_size else 1
max_dynamic_patch = max_dynamic_patch if dynamic_image_size else 1
if use_thumbnail and max_dynamic_patch != 1:
max_dynamic_patch += 1
return min_dynamic_patch, max_dynamic_patch
def get_internvl_target_ratios(
min_num: int,
max_num: int,
) -> list[tuple[int, int]]:
target_ratios = {
(i, j)
for n in range(min_num, max_num + 1)
for i in range(1, n + 1)
for j in range(1, n + 1)
if min_num <= i * j <= max_num
}
return sorted(target_ratios, key=lambda x: x[0] * x[1])
def calculate_internvl_targets(
*,
orig_width: int,
orig_height: int,
target_ratios: list[tuple[int, int]],
image_size: int,
use_thumbnail: bool,
) -> tuple[int, int, int]:
aspect_ratio = orig_width / orig_height
# find the closest aspect ratio to the target
target_aspect_ratio = find_closest_aspect_ratio(
aspect_ratio,
target_ratios,
width=orig_width,
height=orig_height,
image_size=image_size,
)
# calculate the target width and height
target_width = image_size * target_aspect_ratio[0]
target_height = image_size * target_aspect_ratio[1]
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
# add thumbnail image if num_blocks != 1
if use_thumbnail and blocks != 1:
blocks += 1
return blocks, target_width, target_height
# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B
def dynamic_preprocess_internvl(
image: Image.Image,
*,
target_ratios: list[tuple[int, int]],
image_size: int,
use_thumbnail: bool,
) -> list[Image.Image]:
orig_width, orig_height = image.size
# calculate the number of blocks without thumbnail
blocks, target_width, target_height = calculate_internvl_targets(
orig_width=orig_width,
orig_height=orig_height,
target_ratios=target_ratios,
image_size=image_size,
use_thumbnail=False,
)
# resize the image
resized_img = image.resize((target_width, target_height))
processed_images = []
for i in range(blocks):
box = (
(i % (target_width // image_size)) * image_size,
(i // (target_width // image_size)) * image_size,
((i % (target_width // image_size)) + 1) * image_size,
((i // (target_width // image_size)) + 1) * image_size,
)
# split the image
split_img = resized_img.crop(box)
processed_images.append(split_img)
assert len(processed_images) == blocks
if use_thumbnail and len(processed_images) != 1:
thumbnail_img = image.resize((image_size, image_size))
processed_images.append(thumbnail_img)
return processed_images
# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B
def image_to_pixel_values_internvl(
image: Image.Image,
*,
input_size: int,
min_num: int,
max_num: int,
use_thumbnail: bool,
) -> torch.Tensor:
target_ratios = get_internvl_target_ratios(min_num, max_num)
transform = build_transform(input_size=input_size)
images = dynamic_preprocess_internvl(
image,
target_ratios=target_ratios,
image_size=input_size,
use_thumbnail=use_thumbnail,
)
pixel_values = torch.stack([transform(image) for image in images])
return pixel_values
# adapted from https://huggingface.co/OpenGVLab/InternVL2-1B
def video_to_pixel_values_internvl(
video: npt.NDArray,
*,
input_size: int,
min_num: int,
max_num: int,
use_thumbnail: bool,
) -> torch.Tensor:
target_ratios = get_internvl_target_ratios(min_num, max_num)
transform = build_transform(input_size=input_size)
frames_list = list[Image.Image]()
for frame in video:
pil_frame = dynamic_preprocess_internvl(
Image.fromarray(frame, mode="RGB"),
target_ratios=target_ratios,
image_size=input_size,
use_thumbnail=use_thumbnail,
)
assert len(pil_frame) == 1
frames_list.extend(pil_frame)
pixel_values = torch.stack([transform(image) for image in frames_list])
return pixel_values
class InternVLImageProcessor(ImageProcessingMixin):
def __init__(
self,
image_size: int,
min_dynamic_patch: int,
max_dynamic_patch: int,
dynamic_image_size: bool,
use_thumbnail: bool,
) -> None:
self.image_size = image_size
self.min_dynamic_patch = min_dynamic_patch
self.max_dynamic_patch = max_dynamic_patch
self.dynamic_image_size = dynamic_image_size
self.use_thumbnail = use_thumbnail
def resolve_min_max_num(
self,
*,
min_dynamic_patch: int | None = None,
max_dynamic_patch: int | None = None,
dynamic_image_size: bool | None = None,
use_thumbnail: bool | None = None,
) -> tuple[int, int]:
if min_dynamic_patch is None:
min_dynamic_patch = self.min_dynamic_patch
if max_dynamic_patch is None:
max_dynamic_patch = self.max_dynamic_patch
if dynamic_image_size is None:
dynamic_image_size = self.dynamic_image_size
if use_thumbnail is None:
use_thumbnail = self.use_thumbnail
return resolve_internvl_min_max_num(
min_dynamic_patch=min_dynamic_patch,
max_dynamic_patch=max_dynamic_patch,
dynamic_image_size=dynamic_image_size,
use_thumbnail=use_thumbnail,
)
def _images_to_pixel_values_lst(
self,
images: list[Image.Image],
min_dynamic_patch: int | None = None,
max_dynamic_patch: int | None = None,
dynamic_image_size: bool | None = None,
) -> list[torch.Tensor]:
if min_dynamic_patch is None:
min_dynamic_patch = self.min_dynamic_patch
if max_dynamic_patch is None:
max_dynamic_patch = self.max_dynamic_patch
if dynamic_image_size is None:
dynamic_image_size = self.dynamic_image_size
min_num, max_num = resolve_internvl_min_max_num(
min_dynamic_patch=min_dynamic_patch,
max_dynamic_patch=max_dynamic_patch,
dynamic_image_size=dynamic_image_size,
use_thumbnail=False, # Applied in image_to_pixel_values
)
return [
image_to_pixel_values_internvl(
image,
input_size=self.image_size,
min_num=min_num,
max_num=max_num,
use_thumbnail=self.use_thumbnail,
)
for image in images
]
def __call__(
self,
images: Image.Image | list[Image.Image],
*,
min_dynamic_patch: int | None = None,
max_dynamic_patch: int | None = None,
dynamic_image_size: bool | None = None,
return_tensors: str | TensorType | None = None,
**kwargs,
) -> BatchFeature:
images_lst = [images] if not isinstance(images, list) else images
pixel_values_lst = self._images_to_pixel_values_lst(
images_lst,
min_dynamic_patch=min_dynamic_patch,
max_dynamic_patch=max_dynamic_patch,
dynamic_image_size=dynamic_image_size,
)
image_inputs = {
"pixel_values_flat": torch.cat(pixel_values_lst),
"image_num_patches": torch.tensor([len(item) for item in pixel_values_lst]),
}
return BatchFeature(image_inputs, tensor_type=return_tensors)
class InternVLVideoProcessor(BaseVideoProcessor):
def __init__(
self,
image_size: int,
) -> None:
self.image_size = image_size
def _videos_to_pixel_values_lst(
self,
videos: list[npt.NDArray],
) -> list[torch.Tensor]:
return [
video_to_pixel_values_internvl(
video,
input_size=self.image_size,
min_num=1,
max_num=1,
use_thumbnail=False,
)
for video in videos
]
def __call__(
self,
videos: npt.NDArray | list[npt.NDArray],
*,
return_tensors: str | TensorType | None = None,
**kwargs,
) -> BatchFeature:
videos_lst = [videos] if not isinstance(videos, list) else videos
pixel_values_lst = self._videos_to_pixel_values_lst(videos_lst)
image_inputs = {
"pixel_values_flat_video": torch.cat(pixel_values_lst),
"video_num_patches": torch.tensor([len(item) for item in pixel_values_lst]),
}
return BatchFeature(image_inputs, tensor_type=return_tensors)
class InternVLProcessor(ProcessorMixin):
"""This model doesn't define its own HF processor,
so we implement our own one here.
The code to insert image tokens is based on:
https://huggingface.co/OpenGVLab/InternVL2-1B/blob/main/modeling_internvl_chat.py#L252
Code for video processing is adapted from video example:
https://huggingface.co/OpenGVLab/InternVL3-1B#inference-with-transformers
"""
attributes = ["image_processor", "tokenizer", "video_processor"]
def __init__(
self,
image_processor: InternVLImageProcessor,
tokenizer: HfTokenizer,
video_processor: InternVLVideoProcessor | None = None,
*,
image_seq_length: int,
start_image_token: str = "<img>",
end_image_token: str = "</img>",
ctx_image_token: str = "<IMG_CONTEXT>",
ctx_video_token: str | None = None,
) -> None:
self.image_processor = image_processor
self.tokenizer = tokenizer
self.video_processor = video_processor
self.image_seq_length = image_seq_length
self.start_image_token = start_image_token
self.end_image_token = end_image_token
self.ctx_image_token = ctx_image_token
self.ctx_video_token = ctx_video_token
self.start_image_token_id = tokenizer.convert_tokens_to_ids(start_image_token)
self.end_image_token_id = tokenizer.convert_tokens_to_ids(end_image_token)
self.ctx_image_token_id = tokenizer.convert_tokens_to_ids(ctx_image_token)
self.ctx_video_token_id = (
None
if ctx_video_token is None
else tokenizer.convert_tokens_to_ids(ctx_video_token)
)
def resolve_target_ratios(
self,
*,
min_dynamic_patch: int | None = None,
max_dynamic_patch: int | None = None,
dynamic_image_size: bool | None = None,
use_thumbnail: bool | None = None,
) -> list[tuple[int, int]]:
min_num, max_num = self.image_processor.resolve_min_max_num(
min_dynamic_patch=min_dynamic_patch,
max_dynamic_patch=max_dynamic_patch,
dynamic_image_size=dynamic_image_size,
use_thumbnail=use_thumbnail,
)
return get_internvl_target_ratios(min_num, max_num)
def get_num_image_tokens(
self,
*,
image_width: int,
image_height: int,
) -> int:
image_processor = self.image_processor
target_ratios = self.resolve_target_ratios(
use_thumbnail=False, # Applied in calculate_targets
)
num_patches, _, _ = calculate_internvl_targets(
orig_width=image_width,
orig_height=image_height,
image_size=image_processor.image_size,
target_ratios=target_ratios,
use_thumbnail=image_processor.use_thumbnail,
)
return num_patches * self.image_seq_length
def get_image_repl(
self,
num_patches: int | None,
num_features: int | None = None,
) -> PromptUpdateDetails:
if num_patches is None:
assert num_features is not None
else:
num_features = num_patches * self.image_seq_length
repl_features = self.ctx_image_token * num_features
repl_full = self.start_image_token + repl_features + self.end_image_token
full_ids = cached_encode(self.tokenizer, repl_full, add_special_tokens=False)
return PromptUpdateDetails.select_token_id(full_ids, self.ctx_image_token_id)
def get_video_repl(self, num_patches: int) -> PromptUpdateDetails:
assert self.ctx_video_token is not None
repl_features = self.ctx_video_token * self.image_seq_length
repl_features_with_sep = (
self.start_image_token + repl_features + self.end_image_token
)
# num_patches is equal to num_frames
repl_full = "".join(
[f"Frame{i + 1}: {repl_features_with_sep}" for i in range(num_patches)]
)
full_ids = cached_encode(self.tokenizer, repl_full, add_special_tokens=False)
assert self.ctx_video_token_id is not None
return PromptUpdateDetails.select_token_id(full_ids, self.ctx_video_token_id)
def __call__(
self,
text: str | list[str] | None = None,
images: Image.Image | list[Image.Image] | None = None,
videos: npt.NDArray | list[npt.NDArray] | None = None,
*,
min_dynamic_patch: int | None = None,
max_dynamic_patch: int | None = None,
dynamic_image_size: bool | None = None,
return_tensors: str | TensorType | None = None,
**kwargs,
) -> BatchFeature:
if images is not None:
image_inputs = self.image_processor(
images=images,
min_dynamic_patch=min_dynamic_patch,
max_dynamic_patch=max_dynamic_patch,
dynamic_image_size=dynamic_image_size,
return_tensors=return_tensors,
)
image_num_patches = image_inputs["image_num_patches"]
else:
image_inputs = {}
image_num_patches = []
if videos is not None:
if self.video_processor is None:
raise ValueError("This model does not support video inputs")
video_inputs = self.video_processor(
videos=videos,
return_tensors=return_tensors,
)
video_num_patches = video_inputs["video_num_patches"]
else:
video_inputs = {}
video_num_patches = []
if text is not None:
if not isinstance(text, list):
text = [text]
if image_inputs:
image_token = "<image>"
image_index = 0
processed_text = list[str]()
replace_strings = list[str]()
for prompt in text:
new_prompt = prompt
while image_token in new_prompt:
new_prompt = new_prompt.replace(image_token, "<placeholder>", 1)
image_repl = self.get_image_repl(image_num_patches[image_index])
# Convert the token IDs back to text for the
# string-based replacement below
replace_strings.append(
self.tokenizer.decode(
image_repl.full, skip_special_tokens=False
)
)
image_index += 1
while "<placeholder>" in new_prompt:
replace_str = replace_strings.pop(0)
new_prompt = new_prompt.replace("<placeholder>", replace_str, 1)
processed_text.append(new_prompt)
text = processed_text
if video_inputs:
video_token = "<video>"
video_index = 0
processed_text = list[str]()
replace_strings = list[str]()
assert video_token is not None
for prompt in text:
new_prompt = prompt
while video_token in new_prompt:
new_prompt = new_prompt.replace(video_token, "<placeholder>", 1)
video_repl = self.get_video_repl(video_num_patches[video_index])
# Convert the token IDs back to text for the
# string-based replacement below
replace_strings.append(
self.tokenizer.decode(
video_repl.full, skip_special_tokens=False
)
)
video_index += 1
while "<placeholder>" in new_prompt:
replace_str = replace_strings.pop(0)
new_prompt = new_prompt.replace("<placeholder>", replace_str, 1)
processed_text.append(new_prompt)
text = processed_text
text_inputs = self.tokenizer(text, return_tensors=return_tensors)
else:
text_inputs = {}
return BatchFeature(
data={**text_inputs, **image_inputs, **video_inputs},
tensor_type=return_tensors,
)
vllm/transformers_utils/processors/kimi_k25_vision_fused.py
-361
View File
@@ -1,361 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Optimized CPU image processor for Kimi-K2.5/K2.6 vision chunks."""
import io
import json
import math
from typing import Any
import numpy as np
import pybase64 as base64
import torch
from PIL import Image
from transformers.image_processing_utils import BaseImageProcessor, BatchFeature
from transformers.utils import TensorType
from vllm.multimodal.image import open_image
from vllm.utils.import_utils import is_numba_available
from vllm.utils.jit_monitor import numba_workqueue_threading_layer
_IDENTITY_LUT = np.tile(np.arange(256, dtype=np.uint8)[:, None], (1, 3))
if is_numba_available():
from numba import njit, prange
@njit(parallel=True, cache=True)
def _write_fused_patches(
frames: np.ndarray,
out: np.ndarray,
out_offset: int,
new_h: int,
new_w: int,
padded_h: int,
padded_w: int,
patch_size: int,
normalize_lut: np.ndarray,
) -> None:
# frames: [T, new_h, new_w, 3] uint8, without padding.
# out: [total_patches, 3, patch_size, patch_size] float32.
t_size = frames.shape[0]
patch_h = padded_h // patch_size
patch_w = padded_w // patch_size
total = t_size * padded_h * padded_w * 3
hwc = padded_h * padded_w * 3
wc = padded_w * 3
for linear in prange(total):
t = linear // hwc
rem = linear - t * hwc
y = rem // wc
rem = rem - y * wc
x = rem // 3
c = rem - x * 3
value = frames[t, y, x, c] if y < new_h and x < new_w else 0
patch_idx = (
out_offset
+ t * patch_h * patch_w
+ (y // patch_size) * patch_w
+ (x // patch_size)
)
out[patch_idx, c, y % patch_size, x % patch_size] = normalize_lut[value, c]
else:
def _write_fused_patches(*args: Any, **kwargs: Any) -> None:
raise RuntimeError("numba is required for fused Kimi image preprocessing")
def navit_resize_image(
width: int,
height: int,
patch_size: int,
merge_kernel_size: int,
in_patch_limit: int,
patch_limit_on_one_side: int,
fixed_output_tokens: int | None,
) -> dict[str, int]:
s1 = math.sqrt(
in_patch_limit
/ (max(1.0, width // patch_size) * max(1.0, height // patch_size))
)
s2 = patch_limit_on_one_side * patch_size / width
s3 = patch_limit_on_one_side * patch_size / height
scale = min(1.0, s1, s2, s3)
new_w = min(max(1, int(width * scale)), patch_limit_on_one_side * patch_size)
new_h = min(max(1, int(height * scale)), patch_limit_on_one_side * patch_size)
factor = merge_kernel_size * patch_size
pad_height = (factor - new_h % factor) % factor
pad_width = (factor - new_w % factor) % factor
if fixed_output_tokens is not None:
num_tokens = fixed_output_tokens
else:
token_height = (new_h + pad_height) // factor
token_width = (new_w + pad_width) // factor
num_tokens = token_height * token_width
return {
"num_tokens": num_tokens,
"new_width": new_w,
"new_height": new_h,
"pad_width": pad_width,
"pad_height": pad_height,
"sampled_nframes": 1,
}
def navit_resize_video(
width: int,
height: int,
nframes: int,
avg_fps: float,
sample_fps: float,
patch_size: int,
merge_kernel_size: int,
in_patch_limit_each_frame: int,
patch_limit_on_one_side: int,
in_patch_limit_total: int | None,
max_num_frames_each_video: int | None,
fixed_output_tokens_each_frame: int | None,
) -> dict[str, int]:
sample_fps = min(sample_fps, avg_fps)
sampled_nframes = max(round(nframes * sample_fps / avg_fps), 1)
if max_num_frames_each_video is not None:
sampled_nframes = min(sampled_nframes, max_num_frames_each_video)
if in_patch_limit_total is not None:
in_patch_limit_each_frame = min(
round(in_patch_limit_total / sampled_nframes),
in_patch_limit_each_frame,
)
ret = navit_resize_image(
width,
height,
patch_size,
merge_kernel_size,
in_patch_limit_each_frame,
patch_limit_on_one_side,
fixed_output_tokens_each_frame,
)
ret["sampled_nframes"] = sampled_nframes
return ret
def _to_pil(data: Any) -> Image.Image:
if hasattr(data, "media") and hasattr(data, "original_bytes"):
data = data.media
if isinstance(data, Image.Image):
return data if data.mode == "RGB" else data.convert("RGB")
if isinstance(data, str):
if data.startswith("data:"):
raw_base64 = data.split(",", 1)[1]
return open_image(io.BytesIO(base64.b64decode(raw_base64))).convert("RGB")
return open_image(data).convert("RGB")
if isinstance(data, bytes):
return open_image(io.BytesIO(data)).convert("RGB")
raise ValueError(f"Unsupported data type: {type(data)}")
def _ensure_media_type(media: dict[str, Any]) -> dict[str, Any]:
if media["type"] == "image":
media["image"] = _to_pil(media["image"])
return media
if media["type"] == "video_chunk":
media["video_chunk"] = [_to_pil(frame) for frame in media["video_chunk"]]
return media
raise ValueError(f"Unsupported media type: {media['type']}")
class KimiK25FusedVisionProcessor(BaseImageProcessor):
model_type = "kimi_k25"
def __init__(self, media_proc_cfg: dict[str, Any], **kwargs: Any) -> None:
super().__init__(**kwargs)
media_proc_cfg = dict(media_proc_cfg)
merge_kernel_size = media_proc_cfg["merge_kernel_size"]
if isinstance(merge_kernel_size, (list, tuple)):
media_proc_cfg["merge_kernel_size"] = int(merge_kernel_size[0])
self.media_proc_cfg = media_proc_cfg
self.num_frames_per_chunk = media_proc_cfg["temporal_merge_kernel_size"]
values = np.arange(256, dtype=np.float32)[:, None]
image_mean = np.asarray(media_proc_cfg["image_mean"], dtype=np.float32)
image_std_inv = 1.0 / np.asarray(media_proc_cfg["image_std"], dtype=np.float32)
self.normalize_lut = (values / 255.0 - image_mean[None, :]) * image_std_inv[
None, :
]
def media_tokens_calculator(self, media: dict[str, Any]) -> int:
media = _ensure_media_type(media)
ret = self.get_resize_config(media)
return ret["num_tokens"]
def get_resize_config(self, media_input: dict[str, Any]) -> dict[str, int]:
if media_input["type"] == "image":
width, height = media_input["image"].size
return navit_resize_image(
width,
height,
self.media_proc_cfg["patch_size"],
self.media_proc_cfg["merge_kernel_size"],
self.media_proc_cfg["in_patch_limit"],
self.media_proc_cfg["patch_limit_on_one_side"],
self.media_proc_cfg["fixed_output_tokens"],
)
if media_input["type"] == "video_chunk":
frame = media_input["video_chunk"][0]
width, height = frame.size
num_frames = len(media_input["video_chunk"])
in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit_each_frame"]
if in_patch_limit_each_frame is None:
in_patch_limit_each_frame = self.media_proc_cfg["in_patch_limit"]
return navit_resize_video(
width,
height,
num_frames,
1.0,
math.inf,
self.media_proc_cfg["patch_size"],
self.media_proc_cfg["merge_kernel_size"],
in_patch_limit_each_frame,
self.media_proc_cfg["patch_limit_on_one_side"],
self.media_proc_cfg["in_patch_limit_video"],
None,
self.media_proc_cfg["fixed_output_tokens"],
)
raise ValueError(f"Unsupported type: {media_input['type']}")
@staticmethod
def resize_image(image: Image.Image, new_width: int, new_height: int) -> np.ndarray:
image = image.resize((new_width, new_height), resample=Image.Resampling.BICUBIC)
return np.asarray(image)
def preprocess(
self,
medias: list[dict[str, Any]],
return_tensors: str | TensorType | None = None,
do_rescale: bool = True,
do_normalize: bool = True,
) -> BatchFeature:
assert do_rescale == do_normalize
device_normalize = not do_normalize
if not isinstance(medias, list):
medias = [medias]
if not medias:
return BatchFeature(data={}, tensor_type=return_tensors)
if njit is None:
raise RuntimeError("numba is required for fused Kimi image preprocessing")
patch_size = int(self.media_proc_cfg["patch_size"])
prepared = []
grid_thws_np = np.empty((len(medias), 3), dtype=np.int64)
total_patches = 0
for idx, item in enumerate(medias):
item = _ensure_media_type(item)
resize_config = self.get_resize_config(item)
new_width = resize_config["new_width"]
new_height = resize_config["new_height"]
pad_width = resize_config["pad_width"]
pad_height = resize_config["pad_height"]
padded_width = new_width + pad_width
padded_height = new_height + pad_height
if item["type"] == "image":
image_np = self.resize_image(item["image"], new_width, new_height)
frames = image_np[np.newaxis, ...]
elif item["type"] == "video_chunk":
frames = np.stack(
[
self.resize_image(frame, new_width, new_height)
for frame in item["video_chunk"]
],
axis=0,
)
else:
raise ValueError(f"Unsupported type: {item['type']}")
t_size = frames.shape[0]
grid_h = padded_height // patch_size
grid_w = padded_width // patch_size
grid_thws_np[idx, 0] = t_size
grid_thws_np[idx, 1] = grid_h
grid_thws_np[idx, 2] = grid_w
num_patches = t_size * grid_h * grid_w
prepared.append(
(
frames,
new_height,
new_width,
padded_height,
padded_width,
num_patches,
)
)
total_patches += num_patches
lut = _IDENTITY_LUT if device_normalize else self.normalize_lut
pixel_values_np = np.empty(
(total_patches, 3, patch_size, patch_size), dtype=lut.dtype
)
out_offset = 0
with numba_workqueue_threading_layer():
for (
frames,
new_height,
new_width,
padded_height,
padded_width,
num_patches,
) in prepared:
_write_fused_patches(
frames,
pixel_values_np,
out_offset,
new_height,
new_width,
padded_height,
padded_width,
patch_size,
lut,
)
out_offset += num_patches
data = {
"pixel_values": torch.from_numpy(pixel_values_np),
"grid_thws": torch.from_numpy(grid_thws_np),
}
return BatchFeature(data=data, tensor_type=return_tensors)
def __repr__(self):
return f"KimiK25FusedVisionProcessor(media_proc_cfg={self.media_proc_cfg})"
def to_dict(self) -> dict[str, Any]:
output = super().to_dict()
output["media_proc_cfg"] = self.media_proc_cfg
if "media_processor" in output:
del output["media_processor"]
return output
@classmethod
def from_dict(cls, config_dict: dict[str, Any], **kwargs):
config = config_dict.copy()
media_proc_cfg = config.pop("media_proc_cfg", {})
return cls(media_proc_cfg=media_proc_cfg, **config, **kwargs)
def to_json_string(self):
dictionary = self.to_dict()
for key, value in dictionary.items():
if hasattr(value, "tolist"):
dictionary[key] = value.tolist()
return json.dumps(dictionary, indent=2, sort_keys=True) + "\n"
vllm/transformers_utils/processors/muse_glimmer.py
-582
View File
@@ -1,582 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""HuggingFace processor for the MuseGlimmer multimodal model (text + image + video).
This packages MuseGlimmer's vision preprocessing and token-span expansion the standard
HF way, so you can do:
processor = AutoProcessor.from_pretrained(hf_dir, trust_remote_code=True)
messages = [{"role": "user", "content": [
{"type": "image"},
{"type": "text", "text": "Describe this image."},
]}]
text = processor.apply_chat_template(messages, add_generation_prompt=True)
inputs = processor(text=text, images=[pil_image], return_tensors="pt")
out = model.generate(**inputs, max_new_tokens=256)
It reproduces the exact token layout the model expects:
image -> <|image_start|> + <|patch|> * N + <|image_end|>
video -> <|vid_start|>
( "Time: X.Xs" + <|video|> * P [+ <|vid_frame_separator|>] )*
+ <|vid_end|>
where N = floor(H/(patch*ds)) * floor(W/(patch*ds)) (<= 4096) for images and
P = the same grid count per temporal frame-group (<= 144) for video, with the
target (H, W) chosen by MuseGlimmerImageProcessor.compute_image_size /
MuseGlimmerVideoProcessor.compute_video_frame_size (the same grid logic as
modeling_muse_glimmer.MuseGlimmerVisionEncoder).
The chat template emits one sentinel per media item -- ``<|image|>`` for images,
``<|video|>`` for videos -- which ``MuseGlimmerProcessor.__call__`` expands into the
spans above using each media item's computed token count. Vision features are
returned in ``pixel_values`` (a list of per-image / per-frame-group tensors) in
the order the sentinels appear, matching how ``MuseGlimmerModel`` merges them at the
``<|patch|>`` / ``<|video|>`` positions.
"""
from __future__ import annotations
import itertools
import math
from pathlib import Path
import torch
from PIL import Image
from torchvision import transforms as T
from transformers.feature_extraction_utils import BatchFeature
from transformers.image_processing_utils import BaseImageProcessor
from transformers.processing_utils import ProcessorMixin
from transformers.video_processing_utils import BaseVideoProcessor
# Single-token sentinels emitted by the chat template, one per media item.
IMAGE_SENTINEL = "<|image|>" # id 200090; absent from the image/video spans
VIDEO_SENTINEL = "<|video|>" # id 200091; also the per-frame token (see __call__)
# Jinja chat template (multimodal superset of convert_muse_glimmer_to_hf.CHAT_TEMPLATE).
# Message content may be a plain string OR a list of parts
# ({"type": "text"|"image"|"video", ...}); image/video parts render as a single
# sentinel that __call__ expands.
MUSE_GLIMMER_MM_CHAT_TEMPLATE = (
"{{- bos_token -}}"
"{%- macro render_parts(content) -%}"
"{%- if content is string -%}{{- content -}}"
"{%- else -%}"
"{%- for part in content -%}"
"{%- if part['type'] == 'image' -%}{{- '<|image|>' -}}"
"{%- elif part['type'] == 'video' -%}{{- '<|video|>' -}}"
"{%- elif part['type'] == 'text' -%}{{- part['text'] -}}"
"{%- endif -%}"
"{%- endfor -%}"
"{%- endif -%}"
"{%- endmacro -%}"
# At inference (add_generation_prompt), inject a default system message when
# the caller supplied none. Full-transcript rendering
# (add_generation_prompt=False) is left unchanged.
"{%- set ns = namespace(has_system=false) -%}"
"{%- for m in messages -%}"
"{%- if m['role'] == 'system' -%}{%- set ns.has_system = true -%}{%- endif -%}"
"{%- endfor -%}"
"{%- if add_generation_prompt and not ns.has_system -%}"
"{{- '<|start|>system<|message|>You are a helpful assistant.<|eot|>' -}}"
"{%- endif -%}"
"{%- for message in messages -%}"
"{%- set role = message['role'] -%}"
"{%- if role == 'assistant' -%}"
"{%- set recipient = message.get('recipient') -%}"
"{%- set end_turn = message.get('end_turn') -%}"
"{%- if end_turn is none -%}"
"{%- set end_turn = not (recipient and recipient != 'user') -%}"
"{%- endif -%}"
"{{- '<|start|>assistant' -}}"
"{%- if recipient -%}{{- ' to=' + recipient -}}{%- endif -%}"
"{{- '<|message|>' -}}{{- render_parts(message['content']) -}}"
"{{- ('<|eot|>' if end_turn else '<|eom|>') -}}"
"{%- elif role == 'tool' -%}"
"{%- set name = message.get('name', '') -%}"
# Tool content is emitted as-is (string body or interleaved image/text parts).
# Any <tool_output ...> wrapper is baked into the SFT data content; the
# tokenizer does not add it, so the template must not either.
"{{- '<|start|>tool ' + name + '<|message|>' -}}"
"{{- render_parts(message['content']) -}}"
"{{- '<|eot|>' -}}"
"{%- else -%}"
"{%- set header = role -%}"
"{%- if message.get('name') -%}"
"{%- set header = role + ' ' + message['name'] -%}"
"{%- endif -%}"
"{{- '<|start|>' + header + '<|message|>' -}}"
"{{- render_parts(message['content']) -}}"
"{{- '<|eot|>' -}}"
"{%- endif -%}"
"{%- endfor -%}"
"{%- if add_generation_prompt -%}{{- '<|start|>assistant' -}}{%- endif -%}"
)
def _grid_size(
img_w: int, img_h: int, patch_hw: int, max_tokens: int
) -> tuple[int, int, int]:
"""Pick the integer (H, W) grid closest to the aspect ratio under the token cap.
Replicates MuseGlimmerVisionEncoder._compute_grid_size
(modeling_muse_glimmer.py) so the processor needs no torch model import.
Returns (target_h, target_w, n_tokens).
"""
i_nph = img_h / patch_hw
i_npw = img_w / patch_hw
ratio = i_npw / i_nph if i_nph > 0 else 1.0
if i_nph * i_npw > max_tokens:
i_nph = (max_tokens / ratio) ** 0.5
i_npw = i_nph * ratio
candidates = list(
set(
itertools.product(
[math.floor(i_nph), math.ceil(i_nph)],
[math.floor(i_npw), math.ceil(i_npw)],
)
)
)
candidates = [
(nph, npw)
for nph, npw in candidates
if nph >= 1 and npw >= 1 and nph * npw <= max_tokens
]
if not candidates:
candidates = [(max(1, round(i_nph)), max(1, round(i_npw)))]
nph, npw = min(candidates, key=lambda c: abs(c[0] / c[1] - img_h / img_w))
return nph * patch_hw, npw * patch_hw, nph * npw
class MuseGlimmerImageProcessor(BaseImageProcessor):
"""Resize + normalize MuseGlimmer images and compute patch-token counts.
Variable-resolution: each image is resized to the grid that best matches its
aspect ratio under the per-image token cap, then normalized with mean/std 0.5.
Returns per-image tensors (not stacked) because MuseGlimmer consumes a list of
variable-size images. Video frames are handled by MuseGlimmerVideoProcessor.
"""
model_input_names = ["pixel_values"]
def __init__(
self,
patch_size: int = 14,
downsample_factor: int = 2,
max_image_tokens: int = 4096,
image_mean: float = 0.5,
image_std: float = 0.5,
**kwargs,
):
super().__init__(**kwargs)
self.patch_size = patch_size
self.downsample_factor = downsample_factor
self.max_image_tokens = max_image_tokens
self.image_mean = image_mean
self.image_std = image_std
def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor:
# Functional normalize -- no stored Normalize object, so the processor
# stays JSON-serializable for to_dict() / save_pretrained.
return T.functional.normalize(
T.functional.to_tensor(image),
[self.image_mean] * 3,
[self.image_std] * 3,
)
# -- size computation (mirrors modeling_muse_glimmer.MuseGlimmerVisionEncoder) --
def compute_image_size(self, img_w: int, img_h: int) -> tuple[int, int, int]:
ph = self.patch_size * self.downsample_factor
return _grid_size(img_w, img_h, ph, self.max_image_tokens)
# -- preprocessing ---------------------------------------------------------
def preprocess_image(self, image: Image.Image) -> tuple[torch.Tensor, int]:
"""Return (pixel tensor [3, H, W], n_patch_tokens) for one image."""
image = image.convert("RGB")
target_h, target_w, n_tokens = self.compute_image_size(
image.width, image.height
)
image = image.resize((target_w, target_h), Image.LANCZOS)
return self._to_norm_tensor(image), n_tokens
# torchcodec is the training-faithful video decoder; other decoders (torchvision /
# PyAV) seek + color-convert differently and diverge from training. Optional at
# import time so text/image paths load without it; decode_video raises if used.
try:
import torchcodec
except Exception:
torchcodec = None
class MuseGlimmerVideoProcessor(BaseVideoProcessor):
"""MuseGlimmer video preprocessing behind the standard HF
``AutoVideoProcessor`` API.
Exposes MuseGlimmer's training-faithful video handling as a first-class video
processor so you can call ``processor(videos="clip.mp4")`` and downstream
users can discover it via the standard HF ``AutoVideoProcessor`` API plus a
``video_preprocessor_config.json``. It wraps:
* torchcodec decode (matches training; other decoders diverge),
* uniform frame sampling to a whole multiple of ``patch_temporal``,
* real per-group PTS timestamps (rendered as ``Time: X.Xs`` by
MuseGlimmerProcessor),
* ``patch_temporal`` frame-grouping (frames cat on the channel axis ->
``[patch_temporal * 3, H, W]``; the encoder detects video by channel count).
It deliberately overrides ``preprocess`` instead of the BaseVideoProcessor fast
pipeline (group_videos_by_shape / smart_resize / single stacked tensor): MuseGlimmer
consumes a LIST of variable-size group tensors and needs the real per-group PTS,
neither of which the stacked fast path models.
"""
model_input_names = ["pixel_values"]
def __init__(
self,
patch_size: int = 14,
downsample_factor: int = 2,
patch_temporal: int = 2,
max_video_frame_tokens: int = 144,
image_mean: float = 0.5,
image_std: float = 0.5,
video_num_frames: int = 96,
video_sampling_fps: float = 2.0,
**kwargs,
):
super().__init__(**kwargs)
self.patch_size = patch_size
self.downsample_factor = downsample_factor
self.patch_temporal = patch_temporal
self.max_video_frame_tokens = max_video_frame_tokens
self.image_mean = image_mean
self.image_std = image_std
self.video_num_frames = video_num_frames
self.video_sampling_fps = video_sampling_fps
def _to_norm_tensor(self, image: Image.Image) -> torch.Tensor:
return T.functional.normalize(
T.functional.to_tensor(image),
[self.image_mean] * 3,
[self.image_std] * 3,
)
def decode_video(self, video_path: str) -> tuple[list[Image.Image], list[float]]:
"""Sample frames + per-group PTS with torchcodec (the training decode path).
``timestamps[g]`` is the ACTUAL decoded PTS of the first frame in temporal
group ``g``; ``len(frames)`` is a whole multiple of ``patch_temporal``.
"""
if torchcodec is None:
raise RuntimeError(
"torchcodec is required for video decoding (it matches the training "
"decode path). See the Environment Setup steps in hf/README.md."
)
pt = self.patch_temporal
reader = torchcodec.decoders.VideoDecoder(video_path)
total = len(reader)
assert reader.metadata.average_fps is not None, (
f"Video has no FPS metadata: {video_path}"
)
fps = reader.metadata.average_fps
assert self.video_sampling_fps and self.video_sampling_fps > 0, (
f"video_sampling_fps must be positive, got {self.video_sampling_fps}"
)
n = min(
int(total * self.video_sampling_fps / fps), self.video_num_frames, total
)
n = max(pt, (n // pt) * pt)
n = min(n, total)
if n < pt:
raise ValueError(
f"Video has only {total} decodable frame(s) but needs at least "
f"{pt} (one temporal patch): {video_path}"
)
indices = torch.linspace(0, total - 1, n).long().tolist()
frames: list[Image.Image] = []
timestamps: list[float] = []
for j, i in enumerate(indices):
fr = reader[i]
frames.append(
Image.fromarray(fr.data.permute(1, 2, 0).numpy()).convert("RGB")
)
if j % pt == 0:
pts = getattr(fr, "pts_seconds", None)
timestamps.append(float(pts) if pts is not None else i / fps)
return frames, timestamps
def compute_video_frame_size(self, img_w: int, img_h: int) -> tuple[int, int, int]:
"""Pick (H, W, tokens_per_frame) for a video frame under the per-frame cap.
Mirrors MuseGlimmerImageProcessor.compute_image_size but uses the (smaller)
per-frame video token budget; shares the same ``_grid_size`` grid logic as
modeling_muse_glimmer.MuseGlimmerVisionEncoder.
"""
ph = self.patch_size * self.downsample_factor
return _grid_size(img_w, img_h, ph, self.max_video_frame_tokens)
def _group_frames(
self, frames: list[Image.Image]
) -> tuple[list[torch.Tensor], int, int]:
"""Resize/normalize frames and cat ``patch_temporal`` frames per group.
Returns (group_tensors, n_groups, tokens_per_group).
"""
pt = self.patch_temporal
if pt <= 0:
raise ValueError(f"patch_temporal must be positive, got {pt}")
if not frames:
raise ValueError("video must contain at least one frame")
if padding := (-len(frames)) % pt:
frames = [*frames, *([frames[-1]] * padding)]
first = frames[0].convert("RGB")
target_h, target_w, n_tokens = self.compute_video_frame_size(
first.width, first.height
)
groups: list[torch.Tensor] = []
for i in range(0, len(frames), pt):
grp = [
self._to_norm_tensor(
frames[i + j]
.convert("RGB")
.resize((target_w, target_h), Image.LANCZOS)
)
for j in range(pt)
]
groups.append(torch.cat(grp, dim=0))
return groups, len(groups), n_tokens
@staticmethod
def _normalize_videos(videos) -> list:
"""Normalize the ``videos`` arg to a list of per-video items.
A video item is a path (str/Path) or a list of PIL frames. Accepts a
single video (one path, or one list of frames) or a list of those.
"""
if isinstance(videos, (str, Path)):
return [videos]
if videos and not isinstance(videos[0], (list, tuple, str, Path)):
return [videos] # a single list of PIL frames
return list(videos or [])
def preprocess_one(
self,
video: str | Path | list[Image.Image],
timestamps: list[float] | None = None,
) -> tuple[list[torch.Tensor], int, int, list[float]]:
"""Decode (if a path) + group ONE video.
Returns (group_tensors, n_groups, tokens_per_group, group_timestamps). This
is the single per-video implementation shared by ``preprocess`` (the HF
AutoVideoProcessor batch API) and ``MuseGlimmerProcessor.__call__`` (prompt
building), so the two paths cannot drift. For a path, frames + real
per-group PTS come from ``decode_video``; for pre-decoded frames the given
``timestamps`` are used verbatim (empty if not supplied).
"""
if isinstance(video, (str, Path)):
frames, ts = self.decode_video(str(video))
else:
frames = [f.convert("RGB") for f in video]
ts = list(timestamps) if timestamps is not None else []
groups, n_groups, tokens_per_group = self._group_frames(frames)
return groups, n_groups, tokens_per_group, ts
def preprocess(
self,
videos,
video_timestamps: list[list[float]] | None = None,
return_tensors: str | None = "pt",
**kwargs,
) -> BatchFeature:
"""Preprocess video file path(s) or pre-decoded frame list(s).
Accepts a single video (path str, or list of PIL frames) or a list of
those. Returns a BatchFeature carrying, in video order:
* ``pixel_values`` -- flat list of [pt*3, H, W] group tensors,
* ``video_num_groups`` -- groups per video,
* ``video_tokens_per_group`` -- patch-token count per group per video,
* ``video_timestamps`` -- per-group PTS per video,
i.e. exactly what MuseGlimmerProcessor needs to build the
``<|vid_start|> ( Time: X.Xs <|video|>*P [sep] )* <|vid_end|>`` block.
"""
videos = self._normalize_videos(videos)
pixel_values: list[torch.Tensor] = []
num_groups: list[int] = []
tokens_per_group: list[int] = []
out_ts: list[list[float]] = []
for idx, v in enumerate(videos):
ts_in = video_timestamps[idx] if video_timestamps else None
groups, ng, tpg, ts = self.preprocess_one(v, ts_in)
pixel_values += groups
num_groups.append(ng)
tokens_per_group.append(tpg)
out_ts.append(ts)
batch = BatchFeature(
data={
"video_num_groups": num_groups,
"video_tokens_per_group": tokens_per_group,
"video_timestamps": out_ts,
},
tensor_type=None,
)
# Variable-size list -- BatchFeature would fail to stack into one tensor.
batch["pixel_values"] = pixel_values
return batch
class MuseGlimmerProcessor(ProcessorMixin):
"""Bundle MuseGlimmerImageProcessor + MuseGlimmerVideoProcessor + tokenizer;
expand media sentinels into spans.
Images go through ``image_processor`` (MuseGlimmerImageProcessor) and videos through
``video_processor`` (MuseGlimmerVideoProcessor) -- one preprocessing implementation
each, no overlap.
"""
attributes = ["image_processor", "video_processor", "tokenizer"]
image_processor_class = "AutoImageProcessor"
video_processor_class = "AutoVideoProcessor"
tokenizer_class = "TokenizersBackend"
def __init__(
self,
image_processor=None,
video_processor=None,
tokenizer=None,
chat_template=None,
**kwargs,
):
if image_processor is None:
image_processor = MuseGlimmerImageProcessor()
if video_processor is None:
video_processor = MuseGlimmerVideoProcessor()
super().__init__(
image_processor,
video_processor,
tokenizer,
chat_template=chat_template or MUSE_GLIMMER_MM_CHAT_TEMPLATE,
**kwargs,
)
def _sid(self, token: str) -> int:
return self.tokenizer.convert_tokens_to_ids(token)
def _image_block(self, n_tokens: int) -> list[int]:
return (
[self._sid("<|image_start|>")]
+ [self._sid("<|patch|>")] * n_tokens
+ [self._sid("<|image_end|>")]
)
def _video_block(
self,
n_groups: int,
tokens_per_group: int,
timestamps: list[float] | None = None,
) -> list[int]:
"""Per-group ``Time: X.Xs`` + <|video|>*P, separated/terminated.
``timestamps`` (one per temporal group, in seconds) should be the ACTUAL
decoded frame times -- training renders the real factored frame PTS, not a
uniform grid. Falls back to ``g*patch_temporal/fps`` only when timestamps
are not supplied (e.g. callers that don't track frame times); that
approximation can differ from training by the per-frame jitter.
"""
vid = self._sid("<|video|>")
sep = self._sid("<|vid_frame_separator|>")
pt = self.video_processor.patch_temporal
fps = self.video_processor.video_sampling_fps
block = [self._sid("<|vid_start|>")]
for g in range(n_groups):
ts = timestamps[g] if timestamps is not None else g * pt / fps
block += self.tokenizer.encode(f"Time: {ts:.1f}s", add_special_tokens=False)
block += [vid] * tokens_per_group
block.append(sep if g < n_groups - 1 else self._sid("<|vid_end|>"))
return block
def __call__(
self,
text: str | list[str] | None = None,
images: list[Image.Image] | None = None,
videos: list[list[Image.Image] | str] | None = None,
video_timestamps: list[list[float]] | None = None,
return_tensors: str | None = "pt",
**kwargs,
) -> BatchFeature:
if text is None:
raise ValueError(
"`text` is required (use apply_chat_template to build it)."
)
if isinstance(text, (list, tuple)):
if len(text) != 1:
raise ValueError(
"MuseGlimmerProcessor supports a single text sample per call."
)
text = text[0]
images = list(images or [])
videos = list(videos or [])
image_sentinel = self._sid(IMAGE_SENTINEL)
video_sentinel = self._sid(VIDEO_SENTINEL)
# Preprocess media up front so we can expand sentinels in document order.
# Images via image_processor; videos via video_processor (each video item
# is a path -- decoded here -- or a pre-decoded list of PIL frames).
prepped_images = [self.image_processor.preprocess_image(im) for im in images]
prepped_videos = [
self.video_processor.preprocess_one(
v, video_timestamps[i] if video_timestamps else None
)
for i, v in enumerate(videos)
]
ids = self.tokenizer.encode(text, add_special_tokens=False)
n_img = sum(1 for t in ids if t == image_sentinel)
n_vid = sum(1 for t in ids if t == video_sentinel)
if n_img != len(prepped_images):
raise ValueError(
f"{n_img} image sentinel(s) in text but "
f"{len(prepped_images)} image(s) given."
)
if n_vid != len(prepped_videos):
raise ValueError(
f"{n_vid} video sentinel(s) in text but "
f"{len(prepped_videos)} video(s) given."
)
out_ids: list[int] = []
pixel_values: list[torch.Tensor] = []
img_i = vid_i = 0
for tid in ids:
if tid == image_sentinel:
tensor, n_tokens = prepped_images[img_i]
img_i += 1
out_ids += self._image_block(n_tokens)
pixel_values.append(tensor)
elif tid == video_sentinel:
groups, n_groups, tokens_per_group, ts = prepped_videos[vid_i]
vid_i += 1
out_ids += self._video_block(n_groups, tokens_per_group, ts or None)
pixel_values += groups
else:
out_ids.append(tid)
data: dict = {
"input_ids": [out_ids],
"attention_mask": [[1] * len(out_ids)],
}
batch = BatchFeature(data=data, tensor_type=return_tensors)
# Keep pixel_values as a list of variable-size tensors (MuseGlimmerModel
# consumes a list); BatchFeature would fail to stack them into one tensor.
if pixel_values:
batch["pixel_values"] = pixel_values
return batch
__all__ = [
"MuseGlimmerImageProcessor",
"MuseGlimmerVideoProcessor",
"MuseGlimmerProcessor",
"MUSE_GLIMMER_MM_CHAT_TEMPLATE",
]
vllm/transformers_utils/processors/unlimited_ocr.py
-46
View File
@@ -1,46 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Image processor for Unlimited-OCR (baidu/Unlimited-OCR)."""
from PIL import Image
from vllm.logger import init_logger
from vllm.transformers_utils.processors.deepseek_ocr import DeepseekOCRProcessor
logger = init_logger(__name__)
class UnlimitedOCRProcessor(DeepseekOCRProcessor):
"""DeepseekOCRProcessor variant for Unlimited-OCR.
The only behavioural difference from the base processor is a multi-image
safeguard: when more than one image is present, crop ("gundam") mode is
disabled.
Because the effective crop flag then depends on *how many* images are in the
request, the per-item processing output is no longer invariant of sibling
images. ``UnlimitedOCRMultiModalProcessor`` accounts for this by bypassing
the multimodal processing cache for multi-image requests (see its
``_cached_apply_hf_processor``), so the two paths stay consistent.
DeepSeek-OCR does *not* have this restriction because its ``max_crops=6`` is
small enough to be safe for multi-image use.
"""
def tokenize_with_images(
self,
conversation: str,
images: list[Image.Image],
bos: bool = True,
eos: bool = True,
cropping: bool = True,
):
if len(images) > 1 and cropping:
logger.warning_once(
"Unlimited-OCR: crop mode is not supported for multi-image "
"input. Falling back to cropping=False."
)
cropping = False
return super().tokenize_with_images(
conversation, images, bos=bos, eos=eos, cropping=cropping
)
vllm/utils/cpu_triton_utils.py
-444
View File
@@ -1,444 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Contains replacement functions to fallback Triton usages in CPU backend."""
import ctypes
import torch
# For _compute_slot_mapping_kernel in vllm/v1/worker/block_table.py
def _compute_slot_mapping_kernel_impl(
num_tokens: int,
max_num_tokens: int,
query_start_loc: torch.Tensor, # [num_reqs + 1], int32
positions: torch.Tensor, # [num_tokens], int64
block_table: torch.Tensor, # [max_num_reqs, max_num_blocks_per_req], int32
block_table_stride: int, # max_num_blocks_per_req
slot_mapping: torch.Tensor, # [max_num_tokens], int64
KV_CACHE_BLOCK_SIZE: int,
TOTAL_CP_WORLD_SIZE: int = 1,
TOTAL_CP_RANK: int = 0,
CP_KV_CACHE_INTERLEAVE_SIZE: int = 1,
PAD_ID: int = -1,
BLOCK_SIZE: int = 1024,
) -> None:
assert TOTAL_CP_WORLD_SIZE == 1, "Context Parallelism is not supported on CPU."
torch.ops._C.compute_slot_mapping_kernel_impl(
query_start_loc,
positions,
block_table,
slot_mapping,
KV_CACHE_BLOCK_SIZE,
)
def _ensure_int64(t: torch.Tensor) -> torch.Tensor:
return t if t.dtype == torch.int64 else t.to(torch.int64)
def _eagle_prepare_inputs_padded_kernel_impl(
cu_num_draft_tokens,
valid_sampled_tokens_count,
query_start_loc_gpu,
token_indices_to_sample,
num_rejected_tokens_gpu,
num_reqs,
):
# C++ expects int64 for cu_num_draft_tokens, valid_sampled_tokens_count,
# and num_rejected_tokens_gpu, but Python allocates them as int32.
orig_rejected_dtype = num_rejected_tokens_gpu.dtype
rejected_i64 = (
num_rejected_tokens_gpu
if orig_rejected_dtype == torch.int64
else num_rejected_tokens_gpu.to(torch.int64)
)
torch.ops._C.eagle_prepare_inputs_padded_kernel_impl(
_ensure_int64(cu_num_draft_tokens),
_ensure_int64(valid_sampled_tokens_count),
query_start_loc_gpu,
token_indices_to_sample,
rejected_i64,
num_reqs,
)
if orig_rejected_dtype != torch.int64:
num_rejected_tokens_gpu.copy_(rejected_i64.to(orig_rejected_dtype))
def _eagle_prepare_next_token_padded_kernel_impl(
sampled_token_ids,
discard_request_mask,
backup_next_token_ids,
next_token_ids,
valid_sampled_tokens_count,
vocab_size,
num_sampled_tokens_per_req,
num_reqs,
stride=None,
BLOCK_SIZE_TOKENS=None,
):
# C++ reads all integer tensors as int64_t*. Output tensors are written
# in-place so we create int64 copies, call C++, and copy back.
orig_next_dtype = next_token_ids.dtype
orig_valid_dtype = valid_sampled_tokens_count.dtype
next_i64 = _ensure_int64(next_token_ids)
valid_i64 = _ensure_int64(valid_sampled_tokens_count)
torch.ops._C.eagle_prepare_next_token_padded_kernel_impl(
_ensure_int64(sampled_token_ids),
discard_request_mask,
_ensure_int64(backup_next_token_ids),
next_i64,
valid_i64,
vocab_size,
num_sampled_tokens_per_req,
num_reqs,
)
if orig_next_dtype != torch.int64:
next_token_ids.copy_(next_i64.to(orig_next_dtype))
if orig_valid_dtype != torch.int64:
valid_sampled_tokens_count.copy_(valid_i64.to(orig_valid_dtype))
def _eagle_step_slot_mapping_metadata_kernel_impl(
positions,
block_table,
stride,
seq_lens,
out_clamped_positions,
out_slot_mapping,
block_size,
max_model_len,
n_blocks_per_req,
PAD_ID,
batch_size=None,
):
assert batch_size is None or batch_size == positions.shape[0], (
f"batch_size mismatch: {batch_size} vs positions.shape[0]={positions.shape[0]}"
)
torch.ops._C.eagle_step_slot_mapping_metadata_kernel_impl(
positions,
block_table,
seq_lens,
out_clamped_positions,
out_slot_mapping,
block_size,
max_model_len,
PAD_ID,
)
def _copy_and_expand_eagle_inputs_kernel_impl(
target_token_ids_ptr,
target_positions_ptr,
next_token_ids_ptr,
out_input_ids_ptr,
out_positions_ptr,
out_is_rejected_token_mask_ptr,
out_is_masked_token_mask_ptr,
out_new_token_indices_ptr,
out_hidden_state_mapping_ptr,
query_start_loc_ptr,
query_end_loc_ptr,
padding_token_id,
parallel_drafting_token_id,
total_input_tokens,
num_padding_slots_per_request,
shift_input_ids,
BLOCK_SIZE_TOKENS=None,
BLOCK_SIZE_REQS=None,
):
"""Adapter between Triton kernel call convention and C++ implementation.
The Triton kernel uses '_ptr' suffixed parameter names and compile-time
constants (BLOCK_SIZE_TOKENS, BLOCK_SIZE_REQS) which are not needed by
the C++ implementation. C++ reads token id tensors as int64_t*.
Output tensors that are int32 need copy-back after C++ writes int64.
"""
orig_ids_dtype = out_input_ids_ptr.dtype
orig_pos_dtype = out_positions_ptr.dtype
out_ids_i64 = _ensure_int64(out_input_ids_ptr)
out_pos_i64 = _ensure_int64(out_positions_ptr)
torch.ops._C.copy_and_expand_eagle_inputs_kernel_impl(
_ensure_int64(target_token_ids_ptr),
_ensure_int64(target_positions_ptr),
_ensure_int64(next_token_ids_ptr),
out_ids_i64,
out_pos_i64,
out_is_rejected_token_mask_ptr,
out_is_masked_token_mask_ptr,
out_new_token_indices_ptr,
out_hidden_state_mapping_ptr,
query_start_loc_ptr,
query_end_loc_ptr,
padding_token_id,
parallel_drafting_token_id,
total_input_tokens,
num_padding_slots_per_request,
shift_input_ids,
)
if orig_ids_dtype != torch.int64:
out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype))
if orig_pos_dtype != torch.int64:
out_positions_ptr.copy_(out_pos_i64.to(orig_pos_dtype))
def _copy_and_expand_dflash_inputs_kernel_impl(
next_token_ids_ptr,
target_positions_ptr,
out_input_ids_ptr,
out_context_positions_ptr,
out_query_positions_ptr,
out_context_slot_mapping_ptr,
out_query_slot_mapping_ptr,
out_token_indices_ptr,
block_table_ptr,
block_table_stride,
query_start_loc_ptr,
num_rejected_tokens_ptr,
parallel_drafting_token_id,
block_size,
num_query_per_req,
num_speculative_tokens,
total_input_tokens,
BLOCK_SIZE=None,
HAS_NUM_REJECTED=False,
):
"""Adapter between the DFlash Triton launch and the C++ CPU op."""
assert block_table_stride == block_table_ptr.stride(0), (
"block_table_stride mismatch: "
f"{block_table_stride} vs {block_table_ptr.stride(0)}"
)
orig_ids_dtype = out_input_ids_ptr.dtype
orig_context_positions_dtype = out_context_positions_ptr.dtype
orig_query_positions_dtype = out_query_positions_ptr.dtype
orig_context_slot_mapping_dtype = out_context_slot_mapping_ptr.dtype
orig_query_slot_mapping_dtype = out_query_slot_mapping_ptr.dtype
out_ids_i64 = _ensure_int64(out_input_ids_ptr)
out_context_positions_i64 = _ensure_int64(out_context_positions_ptr)
out_query_positions_i64 = _ensure_int64(out_query_positions_ptr)
out_context_slot_mapping_i64 = _ensure_int64(out_context_slot_mapping_ptr)
out_query_slot_mapping_i64 = _ensure_int64(out_query_slot_mapping_ptr)
rejected_i64 = _ensure_int64(num_rejected_tokens_ptr) if HAS_NUM_REJECTED else None
if hasattr(torch.ops._C, "copy_and_expand_dflash_inputs_kernel_impl"):
torch.ops._C.copy_and_expand_dflash_inputs_kernel_impl(
_ensure_int64(next_token_ids_ptr),
_ensure_int64(target_positions_ptr),
out_ids_i64,
out_context_positions_i64,
out_query_positions_i64,
out_context_slot_mapping_i64,
out_query_slot_mapping_i64,
out_token_indices_ptr,
block_table_ptr,
query_start_loc_ptr,
rejected_i64,
parallel_drafting_token_id,
block_size,
num_query_per_req,
num_speculative_tokens,
total_input_tokens,
HAS_NUM_REJECTED,
)
else:
next_ids_i64 = _ensure_int64(next_token_ids_ptr)
target_positions_i64 = _ensure_int64(target_positions_ptr)
block_table_stride = block_table_ptr.stride(0)
num_reqs = query_start_loc_ptr.shape[0] - 1
for req_idx in range(num_reqs):
ctx_start = int(query_start_loc_ptr[req_idx].item())
ctx_end = int(query_start_loc_ptr[req_idx + 1].item())
num_ctx = ctx_end - ctx_start
valid_ctx_end = ctx_end
if rejected_i64 is not None:
valid_ctx_end -= int(rejected_i64[req_idx].item())
# Guard against out-of-bounds: ensure valid_ctx_end > ctx_start.
valid_ctx_end = max(valid_ctx_end, ctx_start + 1)
last_pos = int(target_positions_i64[valid_ctx_end - 1].item())
for j in range(num_ctx):
ctx_idx = ctx_start + j
ctx_pos_idx = min(ctx_idx, total_input_tokens - 1)
position = int(target_positions_i64[ctx_pos_idx].item())
block_num = min(position // block_size, block_table_stride - 1)
block_id = int(block_table_ptr[req_idx, block_num].item())
slot = block_id * block_size + (position % block_size)
out_context_positions_i64[ctx_idx] = position
out_context_slot_mapping_i64[ctx_idx] = slot
for query_off in range(num_query_per_req):
query_out = req_idx * num_query_per_req + query_off
position = last_pos + 1 + query_off
block_num = min(position // block_size, block_table_stride - 1)
block_id = int(block_table_ptr[req_idx, block_num].item())
slot = block_id * block_size + (position % block_size)
out_query_positions_i64[query_out] = position
out_query_slot_mapping_i64[query_out] = slot
out_ids_i64[query_out] = (
int(next_ids_i64[req_idx].item())
if query_off == 0
else parallel_drafting_token_id
)
if query_off > 0:
sample_out_idx = req_idx * num_speculative_tokens + (query_off - 1)
out_token_indices_ptr[sample_out_idx] = query_out
if orig_ids_dtype != torch.int64:
out_input_ids_ptr.copy_(out_ids_i64.to(orig_ids_dtype))
if orig_context_positions_dtype != torch.int64:
out_context_positions_ptr.copy_(
out_context_positions_i64.to(orig_context_positions_dtype)
)
if orig_query_positions_dtype != torch.int64:
out_query_positions_ptr.copy_(
out_query_positions_i64.to(orig_query_positions_dtype)
)
if orig_context_slot_mapping_dtype != torch.int64:
out_context_slot_mapping_ptr.copy_(
out_context_slot_mapping_i64.to(orig_context_slot_mapping_dtype)
)
if orig_query_slot_mapping_dtype != torch.int64:
out_query_slot_mapping_ptr.copy_(
out_query_slot_mapping_i64.to(orig_query_slot_mapping_dtype)
)
def _rejection_greedy_sample_kernel_impl(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
target_argmax,
bonus_token_ids,
is_greedy,
max_spec_len,
uniform_probs=None,
synthetic_conditional_rates=None,
SYNTHETIC_MODE=False,
):
# C++ kernel expects int64 for all integer tensors.
# Note: uniform_probs, synthetic_conditional_rates, and SYNTHETIC_MODE are
# passed by the rejection sampler for synthetic mode support, but are not
# yet implemented in the C++ CPU kernel. We accept them here to maintain
# compatibility with the kernel calling convention.
assert not SYNTHETIC_MODE, "Synthetic acceptance not supported with CPU sampling"
orig_dtype = output_token_ids.dtype
output_token_ids_i64 = _ensure_int64(output_token_ids)
torch.ops._C.rejection_greedy_sample_kernel_impl(
output_token_ids_i64,
_ensure_int64(cu_num_draft_tokens),
_ensure_int64(draft_token_ids),
_ensure_int64(target_argmax),
_ensure_int64(bonus_token_ids),
is_greedy,
max_spec_len,
)
if orig_dtype != torch.int64:
output_token_ids.copy_(output_token_ids_i64.to(orig_dtype))
def _rejection_random_sample_kernel_impl(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
bonus_token_ids,
recovered_token_ids,
uniform_probs,
is_greedy,
max_spec_len,
vocab_size,
synthetic_conditional_rates=None,
NO_DRAFT_PROBS=False,
SYNTHETIC_MODE=False,
):
# C++ kernel expects int64 for all integer tensors and float32 for probs.
# uniform_probs is intentionally float64 in Python to avoid exact-zero
# samples; cast to float32 here for C++ compatibility.
# Note: synthetic_conditional_rates and SYNTHETIC_MODE are passed by the
# rejection sampler for synthetic mode support, but are not yet implemented
# in the C++ CPU kernel. We accept them here to maintain compatibility with
# the kernel calling convention.
assert not SYNTHETIC_MODE, "Synthetic acceptance not supported with CPU sampling"
orig_dtype = output_token_ids.dtype
output_token_ids_i64 = _ensure_int64(output_token_ids)
torch.ops._C.rejection_random_sample_kernel_impl(
output_token_ids_i64,
_ensure_int64(cu_num_draft_tokens),
_ensure_int64(draft_token_ids),
draft_probs,
target_probs,
_ensure_int64(bonus_token_ids),
_ensure_int64(recovered_token_ids),
uniform_probs.to(torch.float32),
is_greedy,
max_spec_len,
vocab_size,
NO_DRAFT_PROBS,
)
if orig_dtype != torch.int64:
output_token_ids.copy_(output_token_ids_i64.to(orig_dtype))
def _expand_kernel_impl(
output,
input_val,
cu_num_tokens,
replace_from,
replace_to,
MAX_NUM_TOKENS=None,
):
torch.ops._C.expand_kernel_impl(
_ensure_int64(output),
_ensure_int64(input_val),
_ensure_int64(cu_num_tokens),
replace_from,
replace_to,
)
def _sample_recovered_tokens_kernel_impl(
output_token_ids,
cu_num_draft_tokens,
draft_token_ids,
draft_probs,
target_probs,
inv_q,
vocab_size,
BLOCK_SIZE=None,
NO_DRAFT_PROBS=False,
USE_FP64_GUMBEL=False,
):
# USE_FP64_GUMBEL only controls the gumbel-noise precision, which the caller
# has already applied to `inv_q` (fp64 vs fp32). The CPU kernel consumes
# `inv_q` directly, so the flag is accepted for interface parity and the
# value is read at its existing dtype.
# C++ reads integer tensors as int64_t*; ensure correct dtype.
orig_dtype = output_token_ids.dtype
output_i64 = _ensure_int64(output_token_ids)
torch.ops._C.sample_recovered_tokens_kernel_impl(
output_i64,
_ensure_int64(cu_num_draft_tokens),
_ensure_int64(draft_token_ids),
draft_probs,
target_probs,
# C++ kernel reads inv_q as float32.
inv_q.to(torch.float32),
vocab_size,
NO_DRAFT_PROBS,
)
if orig_dtype != torch.int64:
output_token_ids.copy_(output_i64.to(orig_dtype))
def _batch_memcpy_impl(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE=None):
# BLOCK_SIZE is unused; kept for signature parity with the Triton kernel.
for src, dst, size in zip(src_ptrs.tolist(), dst_ptrs.tolist(), sizes.tolist()):
ctypes.memmove(dst, src, size)
vllm/v1/attention/backends/composite.py
-510
View File
@@ -1,510 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Reusable attention backend composition and routing policies."""
import functools
import math
from typing import ClassVar, Protocol, cast
import torch
from vllm.config import VllmConfig
from vllm.v1.attention.backend import (
AttentionBackend,
AttentionCGSupport,
AttentionImpl,
AttentionMetadataBuilder,
CommonAttentionMetadata,
)
from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts
from vllm.v1.kv_cache_interface import KVCacheSpec
def requires_mm_prefix(
metadata: CommonAttentionMetadata, *, unclamped_window: bool = False
) -> bool:
ranges = metadata.mm_req_doc_ranges
# With one query per request there are no future keys to unmask. This
# also keeps single-token prefills on the captured FlashInfer path.
if not ranges or (metadata.max_query_len <= 1 and not unclamped_window):
return False
seq_lens = metadata.seq_lens_cpu_upper_bound
assert seq_lens is not None
starts = metadata.query_start_loc_cpu
for req_idx, spans in ranges.items():
query_len = int(starts[req_idx + 1] - starts[req_idx])
if query_len == 0 or (query_len == 1 and not unclamped_window):
continue
end = int(seq_lens[req_idx])
begin = end - query_len
# An unclamped image can also expose keys behind the sliding window.
if unclamped_window and any(
start < stop and start < end and stop >= begin for start, stop in spans
):
return True
# Prefill lengths are exact. Historical image ranges must not divert
# later text/decode queries to the image backend.
if any(
start < stop and start < end - 1 and stop > begin for start, stop in spans
):
return True
return False
def _has_unclamped_window(layers) -> bool:
return any(
getattr(layer, "sliding_window", None) is not None
and not getattr(layer, "mm_prefix_clamp_sliding_window", False)
for layer in layers
)
class CompositeAttentionRouting(Protocol):
"""Select a child backend and declare routing constraints for composition."""
capture_variant: ClassVar[int]
"""Child used during graph capture; replay must select the same child."""
def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None: ...
def select(self, metadata: CommonAttentionMetadata) -> int:
"""Return 0 for the first child or 1 for the second child."""
...
@staticmethod
def get_cudagraph_support(
vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
) -> AttentionCGSupport:
"""Limit graph support to batches whose routing stays invariant."""
...
@staticmethod
def variant_uses_mm_prefix(variant: int) -> bool:
"""Whether this child must support the model's image-mask requirement."""
...
class MMPrefixAttentionRouting(CompositeAttentionRouting):
"""Route image masks to variant zero and causal queries to variant one."""
def __init__(self, layer_names: list[str], vllm_config: VllmConfig) -> None:
layers = vllm_config.compilation_config.static_forward_context
self.unclamped_window = _has_unclamped_window(
layers[name] for name in layer_names
)
def select(self, metadata: CommonAttentionMetadata) -> int:
return int(
not requires_mm_prefix(metadata, unclamped_window=self.unclamped_window)
)
capture_variant: ClassVar[int] = 1
@staticmethod
def get_cudagraph_support(
vllm_config: VllmConfig, kv_cache_spec: KVCacheSpec
) -> AttentionCGSupport:
if _has_unclamped_window(
vllm_config.compilation_config.static_forward_context.values()
):
return AttentionCGSupport.NEVER
return AttentionCGSupport.UNIFORM_SINGLE_TOKEN_DECODE
@staticmethod
def variant_uses_mm_prefix(variant: int) -> bool:
return variant == 0
def _intersect_kernel_block_sizes(
first: type[AttentionBackend],
second: type[AttentionBackend],
kv_cache_spec: KVCacheSpec | None = None,
):
from vllm.v1.attention.backend import MultipleOf
result: list[int | MultipleOf] = []
seen: set[tuple[type, int]] = set()
for lhs in first.get_supported_kernel_block_sizes(kv_cache_spec):
for rhs in second.get_supported_kernel_block_sizes(kv_cache_spec):
candidate: int | MultipleOf | None = None
if isinstance(lhs, MultipleOf) and isinstance(rhs, MultipleOf):
candidate = MultipleOf(math.lcm(lhs.base, rhs.base))
elif isinstance(lhs, MultipleOf) and isinstance(rhs, int):
candidate = rhs if rhs % lhs.base == 0 else None
elif isinstance(lhs, int) and isinstance(rhs, MultipleOf):
candidate = lhs if lhs % rhs.base == 0 else None
elif lhs == rhs:
candidate = lhs
if candidate is None:
continue
value = candidate.base if isinstance(candidate, MultipleOf) else candidate
key = (type(candidate), value)
if key not in seen:
result.append(candidate)
seen.add(key)
return result
@functools.cache
def create_composite_attention_backend(
general_backend: type[AttentionBackend],
causal_backend: type[AttentionBackend],
*,
name: str,
backend_name: str,
module: str,
routing_policy: type[CompositeAttentionRouting],
head_sizes: tuple[int, ...] = (),
kernel_block_sizes: tuple[int, ...] = (),
device_major: int | None = None,
) -> type[AttentionBackend]:
"""Compose two backends with shared storage and an explicit routing policy."""
if general_backend.full_cls_name() == causal_backend.full_cls_name():
return general_backend
general_impl_cls = general_backend.get_impl_cls()
causal_impl_cls = causal_backend.get_impl_cls()
general_builder_cls = general_backend.get_builder_cls()
causal_builder_cls = causal_backend.get_builder_cls()
if (
general_backend.forward_includes_kv_cache_update
or causal_backend.forward_includes_kv_cache_update
):
raise ValueError("Composite backends require separate KV cache updates")
class CompositeAttentionImpl(AttentionImpl):
def __init__(self, *args, **kwargs) -> None:
self.general_impl = cast(AttentionImpl, general_impl_cls(*args, **kwargs))
try:
self.causal_impl = cast(AttentionImpl, causal_impl_cls(*args, **kwargs))
except TypeError as e:
raise ValueError(
f"Causal attention backend {causal_backend.get_name()} "
f"does not accept the attention arguments required by "
f"the general backend {general_backend.get_name()}: {e}"
) from e
for name in (
"num_heads",
"num_kv_heads",
"head_size",
"scale",
"kv_cache_dtype",
):
if hasattr(self.general_impl, name):
setattr(self, name, getattr(self.general_impl, name))
self.supports_quant_query_input = (
self.general_impl.supports_quant_query_input
and self.causal_impl.supports_quant_query_input
)
self.supports_pcp = self.general_impl.supports_pcp
self.supports_dcp = self.causal_impl.supports_dcp
self.can_return_lse_for_decode = self.causal_impl.can_return_lse_for_decode
self.lse_base_on_e = self.causal_impl.lse_base_on_e
self.need_to_return_lse_for_decode = (
self.causal_impl.need_to_return_lse_for_decode
)
self.supports_mtp_with_cp_non_trivial_interleave_size = (
self.causal_impl.supports_mtp_with_cp_non_trivial_interleave_size
)
def get_impl_for_metadata(self, attn_metadata):
return (
self.causal_impl
if getattr(attn_metadata, "_attention_backend_variant", 0)
else self.general_impl
)
def get_impl_variants(self):
return self.general_impl, self.causal_impl
def forward(
self,
layer,
query,
key,
value,
kv_cache,
attn_metadata,
output,
output_scale=None,
output_block_scale=None,
):
impl = self.get_impl_for_metadata(attn_metadata)
return impl.forward(
layer,
query,
key,
value,
kv_cache,
attn_metadata,
output,
output_scale,
output_block_scale,
)
def process_weights_after_loading(self, act_dtype: torch.dtype):
for impl in self.get_impl_variants():
impl.process_weights_after_loading(act_dtype)
def fused_output_quant_supported(self, quant_key):
return all(
impl.fused_output_quant_supported(quant_key)
for impl in self.get_impl_variants()
)
def fused_qk_norm_rope_kvcache_supported(self):
return self.general_impl.fused_qk_norm_rope_kvcache_supported()
def do_qk_norm_rope_kvcache_update(self, *args, **kwargs):
method = self.general_impl.do_qk_norm_rope_kvcache_update
return method(*args, **kwargs)
def fused_rope_kvcache_supported(self):
return self.general_impl.fused_rope_kvcache_supported()
def do_rope_and_kv_cache_update(self, *args, **kwargs):
method = self.general_impl.do_rope_and_kv_cache_update
return method(*args, **kwargs)
def do_kv_cache_update(self, *args, **kwargs):
method = self.general_impl.do_kv_cache_update # type: ignore[attr-defined]
return method(*args, **kwargs)
class CompositeAttentionMetadataBuilder(AttentionMetadataBuilder):
requires_block_table_width = (
general_builder_cls.requires_block_table_width
or causal_builder_cls.requires_block_table_width
)
def __init__(self, kv_cache_spec, layer_names, vllm_config, device, **kwargs):
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
args = (kv_cache_spec, layer_names, vllm_config, device)
# Only the variant that asked for the extra ctor kwargs receives them.
self.general_builder = general_builder_cls(
*args,
**(kwargs if general_builder_cls.requires_block_table_width else {}),
)
self.causal_builder = causal_builder_cls(
*args,
**(kwargs if causal_builder_cls.requires_block_table_width else {}),
)
self.reorder_batch_threshold = min(
(
builder.reorder_batch_threshold
for builder in self._builders
if builder.reorder_batch_threshold is not None
),
default=None,
)
self.routing = routing_policy(layer_names, vllm_config)
@classmethod
def get_cudagraph_support(cls, vllm_config, kv_cache_spec):
return min(
general_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec),
causal_builder_cls.get_cudagraph_support(vllm_config, kv_cache_spec),
routing_policy.get_cudagraph_support(vllm_config, kv_cache_spec),
key=lambda support: support.value,
)
def _builder(self, common_attn_metadata):
backend_variant = self.routing.select(common_attn_metadata)
builder = self.causal_builder if backend_variant else self.general_builder
return builder, backend_variant
def _build(self, method, common_attn_metadata, **kwargs):
builder, backend_variant = self._builder(common_attn_metadata)
metadata = getattr(builder, method)(
common_attn_metadata=common_attn_metadata, **kwargs
)
metadata._attention_backend_variant = backend_variant
return metadata
def build(self, common_prefix_len, common_attn_metadata, fast_build=False):
return self._build(
"build",
common_attn_metadata,
common_prefix_len=common_prefix_len,
fast_build=fast_build,
)
def build_for_cudagraph_capture(self, common_attn_metadata):
variant = routing_policy.capture_variant
metadata = self._builders[variant].build_for_cudagraph_capture(
common_attn_metadata
)
metadata._attention_backend_variant = variant
return metadata
@property
def _builders(self):
return self.general_builder, self.causal_builder
# Forward the optional workspace protocol so that a wrapped builder
# which allocates one (FlashInfer) still joins the runner's cross-group
# sharing instead of allocating a second buffer behind the composite.
def _get_workspace_buffer(self):
for builder in self._builders:
if hasattr(builder, "_get_workspace_buffer"):
workspace = builder._get_workspace_buffer()
# The runner only calls the getter for the group that
# provides the buffer, so hand it to the sibling here.
self.set_workspace_buffer(workspace)
return workspace
return None
def set_workspace_buffer(self, workspace_buffer):
for builder in self._builders:
if hasattr(builder, "set_workspace_buffer"):
builder.set_workspace_buffer(workspace_buffer)
class CompositeAttentionBackend(AttentionBackend):
forward_includes_kv_cache_update = False
general_backend_cls = general_backend
causal_backend_cls = causal_backend
@staticmethod
def get_name():
return backend_name
@staticmethod
def get_impl_cls():
return CompositeAttentionImpl
@staticmethod
def get_builder_cls():
return CompositeAttentionMetadataBuilder
@staticmethod
def get_supported_kernel_block_sizes(
kv_cache_spec: KVCacheSpec | None = None,
):
if kernel_block_sizes:
return [
size
for size in kernel_block_sizes
if all(
backend.supports_block_size(size)
for backend in (general_backend, causal_backend)
)
]
return _intersect_kernel_block_sizes(
general_backend, causal_backend, kv_cache_spec
)
@classmethod
def supports_block_size(cls, block_size):
return bool(
cls.get_supported_kernel_block_sizes()
) and super().supports_block_size(block_size)
@classmethod
def get_supported_head_sizes(cls):
return list(head_sizes)
@classmethod
def supports_head_size(cls, head_size):
return (not head_sizes or head_size in head_sizes) and all(
backend.supports_head_size(head_size)
for backend in (general_backend, causal_backend)
)
@classmethod
def supports_compute_capability(cls, capability):
return (device_major is None or capability.major == device_major) and all(
backend.supports_compute_capability(capability)
for backend in (general_backend, causal_backend)
)
@classmethod
def supports_dtype(cls, dtype):
return all(
backend.supports_dtype(dtype)
for backend in (general_backend, causal_backend)
)
@classmethod
def supports_kv_cache_dtype(cls, dtype):
return all(
backend.supports_kv_cache_dtype(dtype)
for backend in (general_backend, causal_backend)
)
@classmethod
def supports_mm_prefix(cls):
return all(
not routing_policy.variant_uses_mm_prefix(variant)
or backend.supports_mm_prefix()
for variant, backend in enumerate((general_backend, causal_backend))
)
@classmethod
def supports_sliding_window(cls):
return all(
backend.supports_sliding_window()
for backend in (general_backend, causal_backend)
)
@classmethod
def supports_rswa(cls):
return all(
backend.supports_rswa() for backend in (general_backend, causal_backend)
)
@classmethod
def supports_pcp(cls):
return False
@classmethod
def supports_dcp(cls):
return False
@classmethod
def supports_device_cpu_query_lens_mismatch(cls):
return False
@classmethod
def supports_combination(
cls,
head_size,
dtype,
kv_cache_dtype,
block_size,
use_mla,
has_sink,
use_sparse,
use_mm_prefix,
device_capability,
):
for variant, backend in enumerate((general_backend, causal_backend)):
reason = backend.supports_combination(
head_size,
dtype,
kv_cache_dtype,
block_size,
use_mla,
has_sink,
use_sparse,
use_mm_prefix and routing_policy.variant_uses_mm_prefix(variant),
device_capability,
)
if reason:
return f"{backend.get_name()}: {reason}"
return None
@classmethod
def customize_spec(cls, spec):
general_spec = general_backend.customize_spec(spec)
if general_spec != causal_backend.customize_spec(spec):
raise ValueError("Composite backends require the same KV cache spec")
return general_spec
@classmethod
def supported_kv_cache_layouts(cls):
return tuple(
get_supported_kv_cache_layouts((general_backend, causal_backend))
)
CompositeAttentionBackend.__name__ = name
CompositeAttentionBackend.__qualname__ = name
CompositeAttentionBackend.__module__ = module
return CompositeAttentionBackend
vllm/v1/attention/backends/triton_flashinfer.py
-22
View File
@@ -1,22 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Triton image-mask and FlashInfer causal attention composite."""
from vllm.v1.attention.backends.composite import (
MMPrefixAttentionRouting,
create_composite_attention_backend,
)
from vllm.v1.attention.backends.flashinfer import FlashInferBackend
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
TritonFlashInferBackend = create_composite_attention_backend(
TritonAttentionBackend,
FlashInferBackend,
name="TritonFlashInferBackend",
backend_name="TRITON_FLASHINFER",
module=__name__,
routing_policy=MMPrefixAttentionRouting,
head_sizes=(256, 512),
kernel_block_sizes=(64,),
device_major=10,
)
vllm/v1/worker/cpu/__init__.py
View File
vllm/v1/worker/cpu/buffer_utils.py
-14
View File
@@ -1,14 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from collections.abc import Sequence
import torch
class UvaBuffer:
def __init__(self, size: int | Sequence[int], dtype: torch.dtype):
self.cpu = torch.zeros(size, dtype=dtype, device="cpu")
self.np = self.cpu.numpy()
def uva(self, n: int | None = None):
return self.cpu[:n] if n is not None else self.cpu
vllm/v1/worker/cpu/model_runner.py
-16
View File
@@ -1,16 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from vllm.logger import init_logger
from vllm.v1.worker.gpu.model_runner import GPUModelRunner
logger = init_logger(__name__)
class CPUModelRunner(GPUModelRunner):
# TBD: Whether need to move this to Worker?
def warming_up_model(self) -> None:
logger.info("Warming up model for the compilation...")
# Only generate graph for the generic shape
self.profile_run()
logger.info("Warming up done.")
vllm/v1/worker/cpu/shm.py
-104
View File
@@ -1,104 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# isort: skip_file
# ruff: noqa: E402
# mypy: disable-error-code="misc, assignment"
from typing import Any
import numpy as np
# Patch torch APIs
import torch
def noop(*args: Any, **kwargs: Any) -> None:
pass
# Distinct no-op so empty_cache does not alias synchronize: Dynamo's
# handle_synchronize is keyed on that object and asserts on CPU-only hosts.
def empty_cache_noop(*args: Any, **kwargs: Any) -> None:
pass
def fake_pin_memory(self: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
return self
class _EventPlaceholder:
def __init__(self, *args, **kwargs) -> None:
self.record = noop
self.wait = noop
self.synchronize = noop
class _StreamPlaceholder:
def __init__(self, *args, **kwargs) -> None:
self.wait_stream = noop
self.wait_event = noop
self.record_event = noop
self.synchronize = noop
self.query = lambda: True
self.device = torch.device("cpu")
def __enter__(self, *args, **kwargs):
return self
def __exit__(self, exc_type, exc_val, exc_tb):
pass
from vllm.utils.cpu_resource_utils import get_memory_node_info
def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]:
meminfo = get_memory_node_info()
return meminfo.available_memory, meminfo.total_memory
torch.Event = _EventPlaceholder
torch.cuda.Event = _EventPlaceholder
torch.cuda.Stream = _StreamPlaceholder
torch.cuda.set_stream = noop
torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
torch.cuda.stream = lambda *args, **kwargs: _StreamPlaceholder()
torch.accelerator.synchronize = noop
torch.accelerator.empty_cache = empty_cache_noop
torch.Tensor.pin_memory = fake_pin_memory
torch.Tensor.record_stream = noop
torch.accelerator.get_memory_info = get_memory_info
# Patch vLLM torch utils
import vllm.utils.torch_utils as torch_utils
def async_tensor_h2d(
data: list | np.ndarray | torch.Tensor,
device: str | torch.device | None = None,
dtype: torch.dtype | None = None,
out: torch.Tensor | None = None,
) -> torch.Tensor:
if isinstance(data, np.ndarray):
data = torch.from_numpy(data)
if not isinstance(data, torch.Tensor):
data = torch.tensor(data, dtype=dtype, device="cpu")
elif out is None:
return data.to(dtype=dtype)
return data if out is None else out.copy_(data)
torch_utils.async_tensor_h2d = async_tensor_h2d
# Patch model runner APIs
import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils
import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils
gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer
# Patch Triton
from vllm.triton_utils import HAS_TRITON, tl
if HAS_TRITON:
tl.debug_barrier = noop
vllm/v1/worker/cpu_model_runner.py
-188
View File
@@ -1,188 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from contextlib import AbstractContextManager, contextmanager
from typing import Any
import torch
import torch.nn as nn
from vllm.config import (
CompilationMode,
VllmConfig,
)
from vllm.logger import init_logger
from vllm.model_executor.model_loader import get_model
from vllm.tracing import instrument
from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheConfig
from vllm.v1.utils import CpuGpuBuffer
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
logger = init_logger(__name__)
class CPUModelRunner(GPUModelRunner):
def __init__(self, vllm_config: VllmConfig, device: torch.device):
# avoid calling accelerator APIs for methods inherited from super class
_set_torch_accelerator_to_noop()
with _torch_cuda_wrapper():
super().__init__(vllm_config, device)
assert device == torch.device("cpu")
# Note: speculative decoding is now supported on CPU with C++ native impls
self.use_cuda_graph = False
self._postprocess_tensors()
def _postprocess_tensors(self) -> None:
# Note: replace device tensors with cpu tensors
def replace_tensor(obj: Any, cpu_attr_name: str, device_attr_name) -> None:
cpu_tensor = getattr(obj, cpu_attr_name, None)
device_tensor = getattr(obj, device_attr_name, None)
if isinstance(cpu_tensor, torch.Tensor) and isinstance(
device_tensor, torch.Tensor
):
setattr(obj, device_attr_name, cpu_tensor)
for v in vars(self).values():
if isinstance(v, CpuGpuBuffer):
v.gpu = v.cpu
for k, v in vars(self.input_batch).items():
if k.endswith("_cpu_tensor") and isinstance(v, torch.Tensor):
replace_tensor(self.input_batch, k, k[:-11])
for block_table in self.input_batch.block_table.block_tables:
for v in vars(block_table).values():
if isinstance(v, CpuGpuBuffer):
v.gpu = v.cpu
@instrument(span_name="Loading (CPU)")
def load_model(self, load_dummy_weights: bool = False) -> None:
if load_dummy_weights:
raise ValueError(
"Loading dummy weights (needed for elastic EP scale-up) "
"Is not supported by the CPU Model Runner."
)
logger.info("Starting to load model %s...", self.model_config.model)
self.model = get_model(vllm_config=self.vllm_config)
if self.lora_config:
self.model = self.load_lora_model(self.model, self.vllm_config, self.device)
if hasattr(self, "drafter"):
logger.info_once("Loading drafter model...")
self.drafter.load_model(self.model)
self._setup_eagle3_aux_hidden_state_outputs()
def get_model(self) -> nn.Module:
return self.model
@instrument(span_name="Warmup (CPU)")
def warming_up_model(self) -> None:
if self.vllm_config.compilation_config.mode == CompilationMode.NONE:
return
logger.info("Warming up model for the compilation...")
# Only generate graph for the generic shape
with _set_global_compilation_settings(self.vllm_config):
self.profile_run()
logger.info("Warming up done.")
def initialize_kv_cache(
self,
kv_cache_config: KVCacheConfig,
is_profiling: bool = False,
kv_cache_allocation_context: AbstractContextManager | None = None,
) -> None:
super().initialize_kv_cache(
kv_cache_config,
is_profiling,
kv_cache_allocation_context=kv_cache_allocation_context,
)
if self.speculative_config:
if self.speculative_config.use_eagle():
logger.info("EAGLE drafter KV cache initialized for CPU backend")
elif self.speculative_config.uses_draft_model():
logger.info("Draft model KV cache initialized for CPU backend")
def _sync_device(self) -> None:
pass
def _zero_block_ids(self, block_ids: list[int]) -> None:
# Zero full-attention blocks to prevent stale data corruption on partial writes.
# Encoder-only (runner-only) layers are not FullAttentionSpec, so the
# spec filter below already excludes them; no runner-only skip needed.
seen_ptrs: set[int] = set()
for group in self.kv_cache_config.kv_cache_groups:
if not isinstance(group.kv_cache_spec, FullAttentionSpec):
continue
for layer_name in group.layer_names:
ctx = self.compilation_config.static_forward_context.get(layer_name)
if ctx is None:
continue
kv = ctx.kv_cache
if not isinstance(kv, torch.Tensor):
continue
if kv.data_ptr() in seen_ptrs:
continue
seen_ptrs.add(kv.data_ptr())
for block_id in block_ids:
kv[block_id].zero_()
def _to_list(self, sampled_token_ids: torch.Tensor) -> list[list[int]]:
"""CPU-safe version: direct tolist() without CUDA events."""
return sampled_token_ids.tolist()
@contextmanager
def _torch_cuda_wrapper():
class _EventPlaceholder:
def __init__(self, *args, **kwargs) -> None:
self.record = lambda: None
self.synchronize = lambda: None
class _StreamPlaceholder:
def __init__(self, *args, **kwargs) -> None:
self.wait_stream = lambda *a, **kw: None
self.device = torch.device("cpu")
cuda_event = torch.Event
cuda_stream = torch.cuda.Stream
try:
torch.Event = _EventPlaceholder
torch.cuda.Stream = _StreamPlaceholder
yield
finally:
torch.Event = cuda_event
torch.cuda.Stream = cuda_stream
@contextmanager
def _set_global_compilation_settings(config: VllmConfig):
import torch._inductor.config as torch_inductor_config
inductor_config = config.compilation_config.inductor_compile_config
# Note: The MKLDNN and CPPGEMM backend requires freezing parameters.
freezing_value = torch_inductor_config.freezing
try:
if inductor_config.get("max_autotune", False):
torch_inductor_config.freezing = True
yield
finally:
torch_inductor_config.freezing = freezing_value
def _set_torch_accelerator_to_noop() -> None:
def noop(*args: Any, **kwargs: Any) -> None:
pass
# Distinct no-op so empty_cache does not alias synchronize: Dynamo's
# handle_synchronize is keyed on that object and asserts on CPU-only hosts.
def empty_cache_noop(*args: Any, **kwargs: Any) -> None:
pass
torch.accelerator.synchronize = noop
torch.accelerator.empty_cache = empty_cache_noop
vllm/v1/worker/cpu_worker.py
-291
View File
@@ -1,291 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
# Must be imported firstly
import vllm.v1.worker.cpu.shm # noqa # isort: skip
import math
import os
import sys
import psutil
import torch
from vllm import envs
from vllm.config import CompilationMode, VllmConfig
from vllm.logger import init_logger
from vllm.platforms import CpuArchEnum, current_platform
from vllm.utils.cpu_resource_utils import (
get_allowed_cpu_list,
get_memory_node_info,
get_visible_memory_node,
)
from vllm.utils.mem_utils import format_gib
from vllm.utils.torch_utils import set_random_seed
from vllm.v1.worker.cpu_model_runner import CPUModelRunner
from vllm.v1.worker.gpu_worker import Worker, init_worker_distributed_environment
from vllm.v1.worker.worker_base import CompilationTimes
logger = init_logger(__name__)
class CPUWorker(Worker):
def __init__(
self,
vllm_config: VllmConfig,
local_rank: int,
rank: int,
distributed_init_method: str,
is_driver_worker: bool = False,
):
# TODO: use numactl for process setup
# TODO: optimize for `interleaved` policy
# Bind memory node
allowed_memory_nodes = get_visible_memory_node()
allowed_cpu_list = get_allowed_cpu_list()
cpu_core = allowed_cpu_list[0]
# TODO: some CI hosts are not correctly set, change to assertion
# after fix
if cpu_core.numa_node not in allowed_memory_nodes:
logger.warning(
"Node %s is not in available memory nodes %s.",
cpu_core.numa_node,
allowed_memory_nodes,
)
# On s390x, numa_node may be a synthetic book ID that doesn't
# correspond to a real memory node. Fall back to first visible node.
if cpu_core.numa_node in allowed_memory_nodes:
memory_node = cpu_core.numa_node
else:
logger.warning(
"CPU group key %s is not a valid memory node. "
"Falling back to memory node %s.",
cpu_core.numa_node,
allowed_memory_nodes[0],
)
memory_node = allowed_memory_nodes[0]
torch.ops._C.init_cpu_memory_env([memory_node])
memory_status = get_memory_node_info(memory_node)
memory_fraction = vllm_config.cache_config.gpu_memory_utilization
self.requested_cpu_memory = math.ceil(
memory_status.total_memory * memory_fraction
)
available_memory = memory_status.available_memory
if (
vllm_config.cache_config.kv_cache_memory_bytes is None
and self.requested_cpu_memory > available_memory
):
raise ValueError(
f"Available memory on node {cpu_core.numa_node} "
f"({format_gib(available_memory)}/"
f"{format_gib(memory_status.total_memory)} GiB) on startup "
f"is less than desired CPU memory utilization "
f"({vllm_config.cache_config.gpu_memory_utilization}, "
f"{format_gib(self.requested_cpu_memory)} GiB). "
"On the CPU backend, the `--gpu-memory-utilization` flag "
"controls the fraction of CPU memory reserved (despite its "
"name). To resolve: decrease `--gpu-memory-utilization` "
"(e.g. `--gpu-memory-utilization 0.5`) "
"or reduce CPU memory used by other processes."
)
super().__init__(
vllm_config,
local_rank,
rank,
distributed_init_method,
is_driver_worker=is_driver_worker,
)
self.parallel_config.disable_custom_all_reduce = True
def init_device(self):
self.device = torch.device("cpu")
# Check whether critical libraries are loaded
def check_preloaded_libs(name: str) -> bool:
ld_preload_list = os.environ.get("LD_PRELOAD", "")
if name not in ld_preload_list:
logger.warning(
"%s is not found in LD_PRELOAD. "
"For best performance, please follow the section "
"`set LD_PRELOAD` in "
"https://docs.vllm.ai/en/latest/getting_started/installation/cpu/ "
"to setup required pre-loaded libraries.",
name,
)
return False
return True
if sys.platform.startswith("linux"):
check_preloaded_libs("libtcmalloc")
if current_platform.get_cpu_architecture() == CpuArchEnum.X86:
iomp_loaded = check_preloaded_libs("libiomp")
if not iomp_loaded and self.vllm_config.speculative_config is not None:
logger.warning(
"Speculative decoding on CPU without Intel OpenMP in "
"LD_PRELOAD will cause significant performance loss. "
"Please follow the section `set LD_PRELOAD` in "
"https://docs.vllm.ai/en/latest/getting_started/"
"installation/cpu/ "
"to setup libiomp5.",
)
def skip_set_num_threads(x: int):
logger.warning(
"CPU backend doesn't allow to use "
"`torch.set_num_threads` after the thread binding, skip it."
)
torch.set_num_threads = skip_set_num_threads
init_method = self.distributed_init_method
os.environ["VLLM_DIST_IDENT"] = (
os.path.basename(init_method.removeprefix("file://"))
if init_method.startswith("file://")
else init_method.split(":")[-1]
)
# Initialize the distributed environment.
init_worker_distributed_environment(
self.vllm_config,
self.rank,
self.distributed_init_method,
self.local_rank,
current_platform.dist_backend,
)
if self.use_v2_model_runner:
logger.info_once("Using V2 Model Runner")
# Set random seed.
set_random_seed(self.model_config.seed)
# Construct the model runner
if self.use_v2_model_runner:
from vllm.v1.worker.cpu.model_runner import (
CPUModelRunner as CPUModelRunnerV2,
)
self.model_runner: CPUModelRunner = CPUModelRunnerV2( # type: ignore
self.vllm_config, self.device
)
else:
self.model_runner = CPUModelRunner(self.vllm_config, torch.device("cpu"))
def sleep(self, level: int = 1) -> None:
logger.warning("sleep mode is not supported on CPU, ignore it.")
pass
def wake_up(self, tags: list[str] | None = None) -> None:
logger.warning("sleep mode is not supported on CPU, ignore it.")
pass
def discard(self, tags: tuple[str, ...]) -> None:
raise RuntimeError("Selective discard is not supported on CPU.")
def _should_warm_up_model(self) -> bool:
# VLLM_CPU_CI_ENV always skips warmup to save CI time.
if envs.VLLM_CPU_CI_ENV:
return False
# With eager (CompilationMode.NONE) execution and an explicit KV
# cache size, warmup serves no purpose: there's no compiled graph to
# prewarm, and the auto KV-cache-size calc below -- which relies on
# a warmup forward pass having already bumped RSS to a realistic
# steady-state value -- is bypassed whenever the size is explicit.
return not (
self.compilation_config.mode == CompilationMode.NONE
and self.cache_config.kv_cache_memory_bytes is not None
)
def determine_available_memory(self) -> int:
if self._should_warm_up_model():
self.model_runner.warming_up_model()
allowed_cpu_list = get_allowed_cpu_list()
cpu_core = allowed_cpu_list[0]
memory_status = get_memory_node_info(cpu_core.numa_node)
available_memory = memory_status.available_memory
explicit_kv_cache_size = self.cache_config.kv_cache_memory_bytes
kv_cache_size = None
msg = None
if explicit_kv_cache_size is not None:
if explicit_kv_cache_size > available_memory:
raise ValueError(
f"Available memory on node {cpu_core.numa_node} "
f"({format_gib(available_memory)}/"
f"{format_gib(memory_status.total_memory)} GiB) on kv cache"
f" allocation is less than requested memory for kv "
f"({format_gib(explicit_kv_cache_size)} GiB). "
"Decrease --kv-cache-memory-bytes, VLLM_CPU_KVCACHE_SPACE, "
"or reduce CPU memory used by other processes."
)
kv_cache_size = explicit_kv_cache_size
msg = (
f"Explicitly set ({format_gib(kv_cache_size)}/"
f"{format_gib(memory_status.total_memory)}) GiB for KV cache "
f"on node {cpu_core.numa_node}."
)
else:
consumed_memory = psutil.Process(os.getpid()).memory_info().rss
requested_memory_for_kv = int(self.requested_cpu_memory - consumed_memory)
if requested_memory_for_kv <= 0:
# The worker's own footprint already fills its reservation, so
# freeing memory elsewhere cannot help: only a larger
# reservation or an explicit KV cache size can.
in_use = math.ceil(consumed_memory / memory_status.total_memory * 100)
raise ValueError(
f"The worker already uses {format_gib(consumed_memory)} GiB "
"(model weights and runtime), more than the "
f"{format_gib(self.requested_cpu_memory)} GiB of CPU memory "
f"reserved for it on node {cpu_core.numa_node}, leaving none "
"for the KV cache. On the CPU backend, the "
"`--gpu-memory-utilization` flag controls the fraction of CPU "
"memory reserved (despite its name). To resolve: increase "
f"`--gpu-memory-utilization` to more than {in_use / 100:.2f}, "
"or set the KV cache size explicitly with "
"`--kv-cache-memory-bytes` or VLLM_CPU_KVCACHE_SPACE."
)
if requested_memory_for_kv > available_memory:
raise ValueError(
f"Available memory on node {cpu_core.numa_node} "
f"({format_gib(available_memory)}/"
f"{format_gib(memory_status.total_memory)} GiB) on kv cache"
f" allocation is less than requested memory for kv "
f"({format_gib(requested_memory_for_kv)}/"
f"{format_gib(self.requested_cpu_memory)} GiB). "
"Reduce CPU memory used by other processes."
)
kv_cache_size = requested_memory_for_kv
msg = (
f"Auto set ({format_gib(kv_cache_size)}/"
f"{format_gib(memory_status.total_memory)}) GiB for KV cache "
f"on node {cpu_core.numa_node}, with "
f"{format_gib(self.requested_cpu_memory)} GiB requested memory"
f" for the worker. {format_gib(consumed_memory)} GiB"
f" memory was consumed by non-kv usages."
)
logger.info(msg)
return kv_cache_size
def compile_or_warm_up_model(self) -> CompilationTimes:
# Note: the model has been compiled in determine_available_memory(),
# Only compile here for models without kv cache
if len(self.model_runner.kv_caches) == 0 and self._should_warm_up_model():
self.model_runner.warming_up_model()
# Reset the seed to ensure that the random state is not affected by
# the model initialization and profiling.
set_random_seed(
self.model_config.seed, self.parallel_config.data_parallel_index
)
return CompilationTimes(
language_model=self.compilation_config.compilation_time,
encoder=self.compilation_config.encoder_compilation_time,
)