EngineX-Ascend/enginex-ascend-910-vllm
11
0
Fork 1
Code Issues Pull Requests Actions Projects Releases Wiki Activity

init v0.23.0

Browse Source 
Signed-off-by: Sun Ruoxi <sunruoxi@4paradigm.com>
This commit is contained in:
Sun Ruoxi
2026-08-27 15:11:51 +08:00
parent b582a8e7d1
commit 7f8a1b1f7a
2849 changed files with 712887 additions and 22001 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
examples/chat_templates/template_qwen2_audio.jinja Normal file
View File

@@ -0,0 +1,23 @@
{% set audio_count = namespace(value=0) %}
{% for message in messages %}
{% if loop.first and message['role'] != 'system' %}
<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n
{% endif %}
<|im_start|>{{ message['role'] }}\n
{% if message['content'] is string %}
{{ message['content'] }}<|im_end|>\n
{% else %}
{% for content in message['content'] %}
{% if 'audio' in content or 'audio_url' in content or message['type'] == 'audio' or content['type'] == 'audio' %}
{% set audio_count.value = audio_count.value + 1 %}
Audio {{ audio_count.value }}: <|audio_bos|><|AUDIO|><|audio_eos|>\n
{% elif 'text' in content %}
{{ content['text'] }}
{% endif %}
{% endfor %}
<|im_end|>\n
{% endif %}
{% endfor %}
{% if add_generation_prompt %}
<|im_start|>assistant\n
{% endif %}

55
examples/device_print_demo.py Normal file
View File

@@ -0,0 +1,55 @@
import torch
from vllm_ascend.utils import device_print
def compute_and_print(x: torch.Tensor) -> torch.Tensor:
y = torch.square(x) - torch.cos(x)
device_print("device_print from current execution mode")
device_print(7)
device_print(True)
device_print(y)
device_print(f"Compatible with f-strings: {x.dtype = }, {isinstance(x, torch.Tensor) = }")
return y
def main() -> None:
torch.npu.set_device(0)
torch.npu.set_compile_mode(jit_compile=False)
x = torch.arange(1, 28, dtype=torch.float32).reshape(3, 3, 3).npu()
print("=== eager ===", flush=True)
eager_out = compute_and_print(x)
torch.npu.synchronize()
print("=== torch.compile(backend='aot_eager') ===", flush=True)
compiled_compute_and_print = torch.compile(compute_and_print, backend="aot_eager")
compiled_out = compiled_compute_and_print(x)
torch.npu.synchronize()
assert torch.allclose(eager_out, compiled_out), "Outputs from eager and compiled modes do not match."
graph = torch.npu.NPUGraph()
capture_stream = torch.npu.Stream()
x_capture = x.clone()
with torch.npu.stream(capture_stream), torch.npu.graph(graph, stream=capture_stream):
captured_out = compiled_compute_and_print(x_capture)
print("=== replay graph ===", flush=True)
graph.replay()
torch.npu.synchronize()
assert torch.allclose(eager_out, captured_out), "Outputs from eager and graph modes do not match."
print("=== modify input and replay graph ===", flush=True)
x_capture.copy_(torch.arange(28, 1, -1, dtype=torch.float32).reshape(3, 3, 3).npu())
graph.replay()
torch.npu.synchronize()
assert not torch.allclose(eager_out, captured_out), "Outputs from eager and modified graph modes should not match."
if __name__ == "__main__":
main()

206
examples/disaggregated_encoder/disagg_1e1pd_example.sh Normal file
View File

@@ -0,0 +1,206 @@
#!/bin/bash
set -euo pipefail
declare -a PIDS=()
###############################################################################
# Configuration -- override via env before running
###############################################################################
MODEL="${MODEL:-Qwen/Qwen2.5-VL-7B-Instruct}"
LOG_PATH="${LOG_PATH:-./logs}"
mkdir -p $LOG_PATH
ENCODE_PORT="${ENCODE_PORT:-19534}"
PREFILL_DECODE_PORT="${PREFILL_DECODE_PORT:-19535}"
PROXY_PORT="${PROXY_PORT:-10001}"
CARD_E="${CARD_E:-0}"
CARD_PD="${CARD_PD:-1}"
EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/tmp/ec_cache}"
TIMEOUT_SECONDS="${TIMEOUT_SECONDS:-12000}" # wait_for_server timeout
NUM_PROMPTS="${NUM_PROMPTS:-100}" # number of prompts to send in benchmark
###############################################################################
# Helpers
###############################################################################
# Find the git repository root directory
VLLM_ROOT="/vllm-workspace/vllm"
START_TIME=$(date +"%Y%m%d_%H%M%S")
ENC_LOG=$LOG_PATH/encoder_${START_TIME}.log
PD_LOG=$LOG_PATH/pd_${START_TIME}.log
PROXY_LOG=$LOG_PATH/proxy_${START_TIME}.log
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() {
echo "Stopping everything…"
trap - INT TERM USR1 # prevent re-entrancy
# Kill all tracked PIDs
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
echo "Killing process $pid"
kill "$pid" 2>/dev/null
fi
done
# Wait a moment for graceful shutdown
sleep 2
# Force kill any remaining processes
for pid in "${PIDS[@]}"; do
if kill -0 "$pid" 2>/dev/null; then
echo "Force killing process $pid"
kill -9 "$pid" 2>/dev/null
fi
done
# Kill the entire process group as backup
kill -- -$$ 2>/dev/null
echo "All processes stopped."
exit 0
}
trap cleanup INT
trap cleanup USR1
trap cleanup TERM
# clear previous cache
echo "remove previous ec cache folder"
rm -rf $EC_SHARED_STORAGE_PATH
echo "make ec cache folder"
mkdir -p $EC_SHARED_STORAGE_PATH
###############################################################################
# Encoder worker
###############################################################################
ASCEND_RT_VISIBLE_DEVICES="$CARD_E" vllm serve "$MODEL" \
--gpu-memory-utilization 0.01 \
--port "$ENCODE_PORT" \
--enforce-eager \
--enable-request-id-headers \
--no-enable-prefix-caching \
--max-num-batched-tokens 114688 \
--max-num-seqs 128 \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_producer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}' \
>"${ENC_LOG}" 2>&1 &
PIDS+=($!)
###############################################################################
# Prefill+Decode worker
###############################################################################
ASCEND_RT_VISIBLE_DEVICES="$CARD_PD" vllm serve "$MODEL" \
--gpu-memory-utilization 0.9 \
--port "$PREFILL_DECODE_PORT" \
--enforce-eager \
--enable-request-id-headers \
--max-num-seqs 128 \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_consumer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}' \
>"${PD_LOG}" 2>&1 &
PIDS+=($!)
# Wait for workers
wait_for_server $ENCODE_PORT
wait_for_server $PREFILL_DECODE_PORT
###############################################################################
# Proxy
###############################################################################
python ./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" \
>"${PROXY_LOG}" 2>&1 &
PIDS+=($!)
wait_for_server $PROXY_PORT
echo "All services are up!"
###############################################################################
# Single request with local image
###############################################################################
echo "Running single request with local image (non-stream)..."
echo "Running single request with local image (non-stream)..."
base64_image=$(base64 -w 0 "${VLLM_ROOT}/tests/v1/ec_connector/integration/hato.jpg")
cat > /tmp/request.json << EOF
{
"model": "${MODEL}",
"messages": [
{
"role": "system",
"content": "You are a helpful assistant."
},
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {
"url": "data:image/jpg;base64,${base64_image}"
}
},
{
"type": "text",
"text": "What is in this image?"
}
]
}
]
}
EOF
curl http://127.0.0.1:${PROXY_PORT}/v1/chat/completions \
-H "Content-Type: application/json" \
-d @/tmp/request.json
rm -f /tmp/request.json
###############################################################################
# Benchmark
###############################################################################
echo "Running benchmark (stream)..."
vllm bench serve \
--model $MODEL \
--backend openai-chat \
--endpoint /v1/chat/completions \
--dataset-name random-mm \
--seed 0 \
--num-prompts $NUM_PROMPTS \
--port $PROXY_PORT
PIDS+=($!)
# cleanup
echo "cleanup..."
cleanup

749
examples/disaggregated_encoder/disagg_epd_proxy.py Normal file
View File

@@ -0,0 +1,749 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
disagg_encoder_proxy.py
Proxy that routes OpenAI-compatible “/v1/chat/completions” requests to two
clusters:
• encode (multimodal feature extraction)
• decode (language-model inference)
For MM input we:
1. Extract *every* image/audio item.
2. Fire N concurrent requests to the encoder cluster
(one request per item, with **all text removed**).
3. Wait for all of them to succeed.
4. Forward the *original* request to a decode server.
"""
from __future__ import annotations
import argparse
import asyncio
import copy
import logging
import os
import random
import uuid
from collections.abc import AsyncIterator
from enum import Enum
import aiohttp
import uvicorn
from aiohttp import ClientResponse
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
###############################################################################
# FastAPI app & global state
###############################################################################
logging.basicConfig(level=logging.DEBUG, format="%(asctime)s %(levelname)s: %(message)s")
logger = logging.getLogger("proxy")
app = FastAPI()
encode_session: aiohttp.ClientSession | None = None
prefill_session: aiohttp.ClientSession | None = None
decode_session: aiohttp.ClientSession | None = None
###############################################################################
# Utils
###############################################################################
MM_TYPES = {"image_url", "audio_url", "input_audio"}
class EncoderDispatchMode(str, Enum):
SINGLE = "single"
FANOUT = "fanout"
def extract_mm_items(request_data: dict) -> list[dict]:
"""
Return *all* image/audio items that appear anywhere in `messages`.
Each returned dict looks like:
{ "type": "image_url", "image_url": {...} }
"""
items: list[dict] = []
for msg in request_data.get("messages", []):
content = msg.get("content")
if not isinstance(content, list):
continue
for item in content:
if item.get("type") in MM_TYPES:
items.append(item)
return items
async def _encode_fanout(
orig_request: dict,
e_urls: list[str],
req_id: str,
):
logger.info("[%s] Processing multimodal items...", req_id)
mm_items = extract_mm_items(orig_request)
if not mm_items:
logger.info("[%s] No multimodal items, skipping encoder", req_id)
return # nothing to do
logger.info("[%s] got %d multimodal items...", req_id, len(mm_items))
tasks = []
# Round-robin over encode servers to distribute load a bit
url_cycle = (e_urls[i % len(e_urls)] for i in range(len(mm_items)))
for idx, (item, target_url) in enumerate(zip(mm_items, url_cycle)):
# Derive a *child* request id: <parent>:<index>:<random-short>
child_req_id = f"{req_id}:{idx}:{uuid.uuid4().hex[:6]}"
headers = {"x-request-id": child_req_id}
encoder_req = {
# You *may* need to keep additional fields
"model": orig_request.get("model"),
"messages": [
{"role": "user", "content": [item]},
],
# Only need 1 token so the server actually runs the encoder path
"max_tokens": 1,
"stream": False,
}
if encode_session is None:
raise HTTPException(status_code=500, detail="Encode session not initialized")
tasks.append(
encode_session.post(
f"{target_url}/v1/chat/completions",
json=encoder_req,
headers=headers,
)
)
results = await asyncio.gather(*tasks, return_exceptions=True)
# Fail fast if any sub-request failed
for idx, r in enumerate(results):
if isinstance(r, Exception):
logger.error(
"[%s] Encoder request #%d raised exception: %s",
req_id,
idx,
r,
exc_info=r,
)
error_detail = str(r)
if hasattr(r, "status"):
error_detail = f"Status: {r.status}, Error: {error_detail}"
elif hasattr(r, "status_code"):
error_detail = f"Status: {r.status_code}, Error: {error_detail}"
raise HTTPException(status_code=502, detail=f"Encoder request failed: {error_detail}")
if isinstance(r, ClientResponse):
if hasattr(r, "status") and r.status != 200:
try:
detail = await r.text()
except Exception:
detail = "<unable to read body>"
logger.error(
"[%s] Encoder request #%d returned status %s: %s",
req_id,
idx,
r.status,
detail,
)
raise HTTPException(
status_code=r.status,
detail=f"Encoder request failed: {detail}",
)
logger.info("[%s] All %d encoder requests completed successfully", req_id, len(mm_items))
async def _encode_single_request(
orig_request: dict,
e_url: str,
req_id: str,
) -> None:
"""
1. Build one request *per MM item* with all text removed.
2. Send them concurrently to the encode cluster.
3. Raise if any of them fails.
"""
logger.info("[%s] Processing multimodal items...", req_id)
request_data = copy.deepcopy(orig_request)
headers = {"x-request-id": req_id}
request_data["max_tokens"] = 1
request_data["stream"] = False
request_data.pop("stream_options", None)
if "max_completion_tokens" in request_data:
request_data["max_completion_tokens"] = 1
try:
if encode_session is None:
raise HTTPException(status_code=500, detail="Encode session not initialized")
encode_response = await encode_session.post(f"{e_url}/v1/chat/completions", json=request_data, headers=headers)
encode_response.raise_for_status()
if encode_response.status != 200:
encode_text = await encode_response.text()
raise HTTPException(
status_code=encode_response.status,
detail={"error": "Encoder request failed", "message": encode_text},
)
logger.debug("Encoder processing completed successfully for req_id: %s", req_id)
return encode_response
except Exception as e:
logger.error("Encoder processing failed: %s", str(e))
raise HTTPException(
status_code=500,
detail={"error": "Encoder processing error", "message": str(e)},
) from e
logger.info("[%s] Encoder request completed successfully", req_id)
async def fanout_encoder_primer(
orig_request: dict,
req_id: str,
):
mode = app.state.encoder_dispatch_mode
if mode == EncoderDispatchMode.SINGLE:
e_url = random.choice(app.state.e_urls)
await _encode_single_request(orig_request, e_url, req_id)
elif mode == EncoderDispatchMode.FANOUT:
await _encode_fanout(orig_request, app.state.e_urls, req_id)
else:
raise RuntimeError(f"Unknown encoder dispatch mode: {mode}")
async def maybe_prefill(
req_data: dict,
p_url: str,
req_id: str,
) -> dict:
"""
- Do prefill-only task if p_url exist;
- Return modified request data with kv transfer params (for nixl connector)
- Else, skip and return the original request data for decode
"""
if p_url:
logger.info("[%s] Processing through prefill: %s", req_id, p_url)
prefill_response = await process_prefill_stage(req_data, p_url, req_id)
if isinstance(prefill_response, ClientResponse):
# for nixl connector to facilitate kv transfer...
prefill_response_json = await prefill_response.json()
kv_transfer_params = prefill_response_json.get("kv_transfer_params", {})
if kv_transfer_params:
req_data["kv_transfer_params"] = kv_transfer_params
return req_data
else:
return req_data
async def process_prefill_stage(
req_data: dict,
p_url: str,
req_id: str,
) -> ClientResponse:
"""Process request through Prefill stage and return kv_transfer_params"""
logger.info("[%s] Sending prefill request to: %s", req_id, p_url)
prefill_request = req_data.copy()
prefill_request["kv_transfer_params"] = {
"do_remote_decode": True,
"do_remote_prefill": False,
"remote_engine_id": None,
"remote_block_ids": None,
"remote_host": None,
"remote_port": None,
}
prefill_request["stream"] = False
prefill_request["max_tokens"] = 1
if "max_completion_tokens" in prefill_request:
prefill_request["max_completion_tokens"] = 1
if "stream_options" in prefill_request:
del prefill_request["stream_options"]
headers = {"x-request-id": req_id}
try:
if prefill_session is None:
raise HTTPException(status_code=500, detail="Prefill session not initialized")
prefill_response = await prefill_session.post(
f"{p_url}/v1/chat/completions", json=prefill_request, headers=headers
)
prefill_response.raise_for_status()
if prefill_response.status != 200:
error_text = await prefill_response.text()
logger.error(
"[%s] Prefill request failed with status %d: %s",
req_id,
prefill_response.status,
error_text,
)
raise HTTPException(
status_code=prefill_response.status,
detail={"error": "Prefill request failed", "message": error_text},
)
logger.info("[%s] Prefill request completed successfully", req_id)
return prefill_response
except Exception as e:
logger.error("Prefill processing failed: %s", str(e))
raise HTTPException(
status_code=500,
detail={"error": "Prefill processing error", "message": str(e)},
) from e
def has_mm_input(request_data: dict):
if "messages" not in request_data:
return False
for message in request_data["messages"]:
if not isinstance(message.get("content"), list):
continue
for content_item in message["content"]:
if content_item.get("type") in ["image_url", "audio_url", "input_audio"]:
return True
return False
###############################################################################
# Middleware for request/response logging
###############################################################################
@app.middleware("http")
async def log_requests(request: Request, call_next):
"""Middleware to log all incoming requests and responses"""
req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
# Log incoming request
logger.info(
">>> [%s] %s %s from %s",
req_id,
request.method,
request.url.path,
request.client.host if request.client else "unknown",
)
try:
# Process request
response = await call_next(request)
# Log response
logger.info(
"<<< [%s] %s %s completed with status %d",
req_id,
request.method,
request.url.path,
response.status_code,
)
return response
except Exception as e:
# Log errors
logger.exception(
"!!! [%s] %s %s failed with error: %s",
req_id,
request.method,
request.url.path,
str(e),
)
raise
###############################################################################
# FastAPI lifecycle
###############################################################################
@app.on_event("startup")
async def on_startup() -> None:
global encode_session, prefill_session, decode_session
timeout = aiohttp.ClientTimeout(total=100_000)
connector = aiohttp.TCPConnector(limit=0, force_close=False, keepalive_timeout=0)
encode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
if app.state.p_urls:
# only setup if prefill instance(s) exist
prefill_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
decode_session = aiohttp.ClientSession(timeout=timeout, connector=connector)
@app.on_event("shutdown")
async def on_shutdown() -> None:
global encode_session, prefill_session, decode_session
if encode_session:
await encode_session.close()
if prefill_session:
await prefill_session.close()
if decode_session:
await decode_session.close()
###############################################################################
# Core forwarding
###############################################################################
async def forward_non_stream(req_data: dict, req_id: str, p_url: str, d_url: str) -> dict:
try:
# Step 1: Process through Encoder instance (if has MM input)
async def run_encoder():
await fanout_encoder_primer(req_data, req_id)
if has_mm_input(req_data):
await non_stream_retry_wrap(run_encoder)
# Step 2: Process through Prefill instance
async def run_prefill():
return await maybe_prefill(req_data, p_url, req_id)
req_data = await non_stream_retry_wrap(run_prefill)
async def run_decode_non_stream():
# Step 3: Process through Decode instance
logger.info("[%s] Forwarding to decode: %s", req_id, d_url)
headers = {"x-request-id": req_id}
# Non-streaming response
if decode_session is None:
raise HTTPException(status_code=500, detail="Decode session not initialized")
async with decode_session.post(f"{d_url}/v1/chat/completions", json=req_data, headers=headers) as resp:
resp.raise_for_status()
return await resp.json()
return await non_stream_retry_wrap(run_decode_non_stream)
except HTTPException:
raise
except Exception as e:
logger.exception("[%s] Error in forward_non_stream: %s", req_id, str(e))
raise HTTPException(status_code=500, detail=f"Proxy error: {str(e)}") from e
async def stream_retry_wrap(forward_func, max_retries: int = 3, delay: float = 0.001):
last_exc = None
first_chunk_sent = False
for attempt in range(max_retries):
try:
async for chunk in forward_func():
first_chunk_sent = True
yield chunk
return
except Exception as e:
if first_chunk_sent:
raise
if isinstance(e, HTTPException) and e.status_code < 500:
raise
last_exc = e
logger.warning(
"attempt %s / %s failed retrying... ",
attempt + 1,
max_retries,
)
await asyncio.sleep(delay * (attempt + 1))
raise RuntimeError(f"all {max_retries} retries failed.") from last_exc
async def non_stream_retry_wrap(forward_func, max_retries: int = 3, delay: float = 0.001):
last_exc = None
for attempt in range(max_retries):
try:
result = await forward_func()
return result
except Exception as e:
if isinstance(e, HTTPException) and e.status_code < 500:
raise
last_exc = e
logger.warning(
"attempt %s / %s failed retrying... ",
attempt + 1,
max_retries,
)
await asyncio.sleep(delay * (attempt + 1))
raise RuntimeError(f"all {max_retries} retries failed.") from last_exc
async def forward_stream(req_data: dict, req_id: str, p_url: str, d_url: str) -> AsyncIterator[str]:
try:
# Step 1: Process through Encoder instance (if has MM input)
async def run_encoder():
await fanout_encoder_primer(req_data, req_id)
if has_mm_input(req_data):
await non_stream_retry_wrap(run_encoder)
# Step 2: Process through Prefill instance
async def run_prefill():
return await maybe_prefill(req_data, p_url, req_id)
req_data = await non_stream_retry_wrap(run_prefill)
async def run_decode_stream():
# Step 3: Process through Decode instance
logger.info("[%s] Starting streaming from decode: %s", req_id, d_url)
headers = {"x-request-id": req_id}
# Streaming response
if decode_session is None:
raise HTTPException(status_code=500, detail="Decode session not initialized")
async with decode_session.post(
f"{d_url}/v1/chat/completions",
json=req_data,
headers=headers,
) as resp:
resp.raise_for_status()
async for chunk in resp.content.iter_chunked(1024):
if chunk:
yield chunk.decode("utf-8", errors="ignore")
logger.info("[%s] Streaming completed", req_id)
async for chunk in stream_retry_wrap(run_decode_stream):
yield chunk
except HTTPException:
logger.exception("[%s] HTTPException in forward_stream", req_id)
raise
except Exception as e:
logger.exception("[%s] Error in forward_stream: %s", req_id, str(e))
raise HTTPException(status_code=500, detail=f"Proxy streaming error: {str(e)}") from e
###############################################################################
# Public routes
###############################################################################
@app.post("/v1/chat/completions")
async def chat_completions(request: Request):
try:
req_data = await request.json()
req_id = request.headers.get("x-request-id", str(uuid.uuid4()))
p_url = random.choice(app.state.p_urls) if app.state.p_urls else None
d_url = random.choice(app.state.d_urls)
is_streaming = req_data.get("stream", False)
if is_streaming:
return StreamingResponse(
forward_stream(req_data, req_id, p_url, d_url),
media_type="text/event-stream",
)
result = await forward_non_stream(req_data, req_id, p_url, d_url)
return JSONResponse(content=result)
except HTTPException:
raise
except Exception as e:
logger.exception("Error in chat_completions endpoint: %s", str(e))
raise HTTPException(status_code=500, detail=f"Request processing error: {str(e)}") from e
@app.get("/v1/models")
async def list_models():
if decode_session is None:
raise HTTPException(status_code=500, detail="Decode session not initialized")
async with decode_session.get(f"{app.state.d_urls[0]}/v1/models") as resp:
resp.raise_for_status()
return await resp.json()
@app.get("/health")
async def health_check():
async def healthy(urls, session):
if not urls:
return "empty"
for u in urls:
try:
if session is None:
return "unhealthy"
async with session.get(f"{u}/health") as resp:
resp.raise_for_status()
except Exception:
return "unhealthy"
return "healthy"
e_status, p_status, d_status = await asyncio.gather(
healthy(app.state.e_urls, encode_session),
healthy(app.state.p_urls, prefill_session),
healthy(app.state.d_urls, decode_session),
)
overall_healthy = all(status != "unhealthy" for status in (e_status, p_status, d_status))
status_code = 200 if overall_healthy else 503
return JSONResponse(
{
"proxy": "healthy",
"encode_cluster": e_status,
"prefill_cluster": p_status,
"decode_cluster": d_status,
},
status_code=status_code,
)
###############################################################################
# Simple profiler fan-out (unchanged except for sessions)
###############################################################################
async def _post_if_available(
session: aiohttp.ClientSession,
url: str,
payload: dict,
headers: dict,
) -> dict | None:
"""
POST `payload` to `url`.
Returns
-------
• The decoded JSON body on success (2xx)
• None if the endpoint does not exist (404)
• Raises for anything else.
"""
try:
if session is None:
return None
resp = await session.post(url, json=payload, headers=headers)
if resp.status == 404: # profiling disabled on that server
logger.warning("Profiling endpoint missing on %s", url)
return None
resp.raise_for_status()
return await resp.json(content_type=None)
except aiohttp.ClientResponseError as exc:
# Pass 404 through the branch above, re-raise everything else
if exc.status == 404:
logger.warning("Profiling endpoint missing on %s", url)
return None
raise
except Exception:
# Network errors etc.: propagate
raise
async def _profile_cmd(cmd: str, payload: dict, e_url: str, p_url: str, d_url: str):
"""
Fire & forget to both clusters, tolerate 404.
"""
headers = {"Authorization": f"Bearer {os.getenv('OPENAI_API_KEY', '')}"}
encode_task = _post_if_available(encode_session, f"{e_url}/{cmd}_profile", payload, headers)
prefill_task = (
_post_if_available(prefill_session, f"{p_url}/{cmd}_profile", payload, headers)
if p_url is not None
else asyncio.sleep(0)
)
decode_task = _post_if_available(decode_session, f"{d_url}/{cmd}_profile", payload, headers)
encode_res, prefill_res, decode_res = await asyncio.gather(encode_task, prefill_task, decode_task)
# If *all* clusters said “I don’t have that route”, surface an error
if encode_res is prefill_res is decode_res is None:
raise HTTPException(
status_code=503,
detail="Profiling endpoints are disabled on all clusters",
)
return {
"encode": encode_res, # may be None
"prefill": prefill_res, # may be None
"decode": decode_res, # may be None
}
@app.post("/start_profile")
async def start_profile(request: Request):
body = await request.json()
# TODO: handle multi urls properly
e_url = random.choice(app.state.e_urls)
p_url = random.choice(app.state.p_urls) if app.state.p_urls else None
d_url = random.choice(app.state.d_urls)
return await _profile_cmd("start", body, e_url, p_url, d_url)
@app.post("/stop_profile")
async def stop_profile(request: Request):
body = await request.json()
# TODO: handle multi urls properly
e_url = random.choice(app.state.e_urls)
p_url = random.choice(app.state.p_urls) if app.state.p_urls else None
d_url = random.choice(app.state.d_urls)
return await _profile_cmd("stop", body, e_url, p_url, d_url)
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--host", default="0.0.0.0")
parser.add_argument("--port", type=int, default=8000)
parser.add_argument(
"--encode-servers-urls",
required=True,
help='Comma-separated encode URLs ("http://e1:8001,http://e2:8001")',
)
parser.add_argument(
"--prefill-servers-urls",
required=True,
help='Comma-separated prefill URLs ("http://p1:8003,http://p2:8004") to enable E->P->D, '
'set "disable" or "none" to enable E->PD',
)
parser.add_argument(
"--decode-servers-urls",
required=True,
help='Comma-separated decode URLs ("http://d1:8005,http://d2:8006")',
)
parser.add_argument(
"--encoder-dispatch-mode",
choices=["single", "fanout"],
default="single",
help="Encoder dispatch mode: single (one request) or fanout (per-MM-item)",
)
args = parser.parse_args()
app.state.e_urls = [u.strip() for u in args.encode_servers_urls.split(",") if u.strip()]
app.state.d_urls = [u.strip() for u in args.decode_servers_urls.split(",") if u.strip()]
# handle prefill instances
if args.prefill_servers_urls.lower() in ("disable", "none", ""):
app.state.p_urls = []
logger.info("Disaggregated prefill phase explicitly disabled by user. Running E + PD...")
else:
app.state.p_urls = [u.strip() for u in args.prefill_servers_urls.split(",") if u.strip()]
logger.info("Disaggregated prefill phase is enabled. Running E + P + D...")
app.state.encoder_dispatch_mode = EncoderDispatchMode(args.encoder_dispatch_mode)
logger.info("Proxy listening on %s:%s", args.host, args.port)
logger.info("Encode servers: %s", app.state.e_urls)
logger.info("Prefill instances %s", app.state.p_urls)
logger.info("Decode servers: %s", app.state.d_urls)
uvicorn.run(
app,
host=args.host,
port=args.port,
log_level="info",
loop="uvloop",
access_log=True,
)

640
examples/disaggregated_prefill_v1/load_balance_proxy_layerwise_server_example.py Normal file
View File

@@ -0,0 +1,640 @@
# Adapted from https://github.com/vllm-project/vllm/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
# SPDX-License-Identifier: Apache-2.0
#
# Tutorial: Using the Load Balance Proxy Server Example
#
# This proxy server is designed to distribute requests between multiple
# "prefiller" and "decoder" backend servers for large language model inference.
# It is useful for scaling out inference workloads and balancing load across
# multiple backend instances.
#
# Features:
# - Load balances requests to multiple prefiller and decoder servers.
# - Supports OpenAI-compatible /v1/completions and /v1/chat/completions endpoints.
# - Streams responses from backend servers to clients.
#
# Prerequisites:
# - Python 3.8+
# - Install dependencies:
# pip install fastapi<0.124.0 httpx uvicorn vllm
#
# Step 1: Start Your Backend Servers
# ----------------------------------
# You need to have at least one prefiller and one decoder backend running.
# These can be mock servers or actual vLLM servers.
#
# For testing, you can use the provided mock server:
#
# vllm serve --host 0.0.0.0 --port 8100 ... # Prefiller 1
# vllm serve --host 0.0.0.0 --port 8101 ... # Prefiller 2
# vllm serve --host 0.0.0.0 --port 8200 ... # Decoder 1
# vllm serve --host 0.0.0.0 --port 8201 ... # Decoder 2
#
# Step 2: Start the Proxy Server
# ------------------------------
# Run the proxy server, specifying the host/port for each prefiller and decoder:
#
# python load_balance_proxy_server_example.py \
# --host 0.0.0.0 --port 9000 \
# --prefiller-hosts 127.0.0.1 127.0.0.1 \
# --prefiller-ports 8100 8101 \
# --decoder-hosts 127.0.0.1 127.0.0.1 \
# --decoder-ports 8200 8201
#
# This will start the proxy on port 9000, load balancing between two prefiller
# and two decoder servers.
#
# Step 3: Send a Request to the Proxy
# -----------------------------------
# You can now send OpenAI-compatible requests to the proxy. For example:
#
# curl -X POST http://localhost:9000/v1/completions \
# -H "Content-Type: application/json" \
# -d '{
# "model": "your-model",
# "prompt": "The quick brown fox jumps over the lazy dog",
# "max_tokens": 16
# }'
#
# Or for chat completions:
#
# curl -X POST http://localhost:9000/v1/chat/completions \
# -H "Content-Type: application/json" \
# -d '{
# "model": "your-model",
# "messages": [{"role": "user", "content": "Hello!"}],
# "max_tokens": 16
# }'
#
# Step 4: Health Check
# --------------------
# To check if the proxy is running and see how many backend instances are
# connected, use:
#
# curl http://localhost:9000/healthcheck
#
# This will return a JSON object with the status and the number of prefiller
# and decoder instances.
#
# Notes:
# - You can scale the number of prefiller and decoder servers as needed.
# - The proxy will round-robin requests to balance load.
# - For production, ensure your backend servers are robust and secure.
#
# For more details, see the code and comments in this file.
import argparse
import asyncio
import copy
import functools
import heapq
import ipaddress
import json
import os
import sys
import uuid
from contextlib import asynccontextmanager
import httpx
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from vllm.logger import init_logger
logger = init_logger(__name__)
# Add uvloop for faster event loop if available
try:
import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
except ImportError:
pass
class ServerState:
def __init__(self, host, port):
self.host = host
self.port = port
self.url = f"http://{host}:{port}/v1"
# Auto-completion for ipv6
try:
ip = ipaddress.ip_address(self.host)
if isinstance(ip, ipaddress.IPv6Address):
self.url = f"http://[{host}]:{port}/v1"
except Exception:
pass
self.client = httpx.AsyncClient(
timeout=None,
base_url=self.url,
limits=httpx.Limits(max_connections=100000, max_keepalive_connections=100000),
)
self.active_tokens = 0
self.active_kv_cache = 0 # Only for prefiller
self.active_requests = 0 # Number of active requests
self.aborted_requests = set() # Track aborted requests
# Removed individual server lock - will use global locks instead
class ProxyState:
def __init__(self, prefiller_instances, decoder_instances):
self.prefillers: list[ServerState] = [ServerState(h, p) for h, p in prefiller_instances]
self.decoders: list[ServerState] = [ServerState(h, p) for h, p in decoder_instances]
self.req_to_prefiller = {}
self.req_id_lock = asyncio.Lock()
# Removed selection locks - no longer needed for synchronous methods
# Initialize priority queues for efficient server selection
# Each entry is (priority_score, server_index, server_reference)
# Lower priority score = higher priority (less loaded)
self.prefiller_heap = [(0, i, server) for i, server in enumerate(self.prefillers)]
self.decoder_heap = [(0, i, server) for i, server in enumerate(self.decoders)]
heapq.heapify(self.prefiller_heap)
heapq.heapify(self.decoder_heap)
self.req_id_future = {}
self.req_data_dict = {}
def _update_prefiller_priority(self, server_idx: int):
"""Update the priority of a prefiller server in the heap."""
server = self.prefillers[server_idx]
# Priority based on active_tokens and active_kv_cache
priority = server.active_tokens + server.active_kv_cache * 0.3
# Remove old entry and add new one
self.prefiller_heap = [(p, i, s) for p, i, s in self.prefiller_heap if i != server_idx]
heapq.heappush(self.prefiller_heap, (priority, server_idx, server)) # type: ignore
def _update_decoder_priority(self, server_idx: int):
"""Update the priority of a decoder server in the heap."""
server = self.decoders[server_idx]
priority = server.active_tokens
# Remove old entry and add new one
self.decoder_heap = [(p, i, s) for p, i, s in self.decoder_heap if i != server_idx]
heapq.heappush(self.decoder_heap, (priority, server_idx, server)) # type: ignore
def abort_prefiller_request(self, server_idx: int, request_id): # Changed to synchronous
"""
Mark a request as aborted. This will helps to release kv cache in
prefiller node.
"""
# No lock needed - atomic operation
self.prefillers[server_idx].aborted_requests.add(request_id)
def acquire_aborted_prefiller_requests(self, server_idx: int): # Changed to synchronous
"""
Get the set of aborted requests and clear it.
This is used to release kv cache in prefiller node.
"""
# No lock needed - atomic operation
aborted_requests = self.prefillers[server_idx].aborted_requests.copy()
self.prefillers[server_idx].aborted_requests.clear()
return aborted_requests
async def next_req_id(self):
async with self.req_id_lock:
return str(uuid.uuid4())
def select_prefiller(self, token_count): # Changed to synchronous
# No lock needed - entire function is atomic
if not self.prefiller_heap:
raise RuntimeError("No prefiller servers available")
priority, chosen, server = heapq.heappop(self.prefiller_heap)
# Update the chosen server atomically
self.prefillers[chosen].active_tokens += token_count
self.prefillers[chosen].active_kv_cache += token_count
# Update priority and re-add to heap
self._update_prefiller_priority(chosen)
return chosen
def release_prefiller(self, idx, token_count): # Changed to synchronous
# No lock needed - atomic operation
self.prefillers[idx].active_tokens -= token_count
# Update priority queue after releasing
self._update_prefiller_priority(idx)
def release_prefiller_kv(self, idx, token_count): # Changed to synchronous
# No lock needed - atomic operation
if self.prefillers[idx].active_kv_cache > 0:
self.prefillers[idx].active_kv_cache -= token_count
# Update priority queue after releasing
self._update_prefiller_priority(idx)
def select_decoder(self, token_count): # Changed to synchronous
# No lock needed - entire function is atomic
if not self.decoder_heap:
raise RuntimeError("No decoder servers available")
priority, chosen, server = heapq.heappop(self.decoder_heap)
# Update the chosen server atomically
self.decoders[chosen].active_tokens += token_count
# Update priority and re-add to heap
self._update_decoder_priority(chosen)
return chosen
def release_decoder(self, idx, token_count): # Changed to synchronous
# No lock needed - atomic operation
self.decoders[idx].active_tokens -= token_count
# Update priority queue after releasing
self._update_decoder_priority(idx)
# Omni_infer's calculate_input_scores function
def calculate_prefill_scores(self, request_length: int) -> float:
length_score = request_length / 4.0
input_score = length_score * 0.0345 + 120.0745
return input_score
def calculate_decode_scores(self, request_length: int) -> float:
return request_length
proxy_state = None
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--host", type=str, default="localhost")
parser.add_argument("--prefiller-hosts", type=str, nargs="+", default=["localhost"])
parser.add_argument("--prefiller-ports", type=int, nargs="+", default=[8001])
parser.add_argument("--decoder-hosts", type=str, nargs="+", default=["localhost"])
parser.add_argument("--decoder-ports", type=int, nargs="+", default=[8002])
parser.add_argument("--max-retries", type=int, default=3, help="Maximum number of retries for HTTP requests")
parser.add_argument(
"--retry-delay", type=float, default=0.001, help="Base delay (seconds) for exponential backoff retries"
)
args = parser.parse_args()
logger.info(
"Decoder hosts will access Proxy host:port/metaserver, ensure that %s can access %s:%s/metaserver",
set(args.decoder_hosts),
args.host,
args.port,
)
# Wildcard address is not allowed for layerwise connector
if args.host in ["0.0.0.0", "::", "0:0:0:0:0:0:0:0"]:
raise ValueError(
f"Decoder hosts will access Proxy host:port/metaserver, to avoid configuration errors, "
f"the Wildcard Address {args.host} is not allowed for proxy"
)
if len(args.prefiller_hosts) != len(args.prefiller_ports):
raise ValueError("Number of prefiller hosts must match number of prefiller ports")
if len(args.decoder_hosts) != len(args.decoder_ports):
raise ValueError("Number of decoder hosts must match number of decoder ports")
args.prefiller_instances = list(zip(args.prefiller_hosts, args.prefiller_ports))
args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
return args
@asynccontextmanager
async def lifespan(app: FastAPI):
global proxy_state
proxy_state = ProxyState(global_args.prefiller_instances, global_args.decoder_instances)
print(f"Initialized {len(proxy_state.prefillers)} prefill clients and {len(proxy_state.decoders)} decode clients.")
yield
for p in proxy_state.prefillers:
await p.client.aclose()
for d in proxy_state.decoders:
await d.client.aclose()
async def listen_for_disconnect(request: Request) -> None:
"""Return if a disconnect message is received"""
while True:
message = await request.receive()
if message["type"] == "http.disconnect":
break
def with_cancellation(handler_func):
@functools.wraps(handler_func)
async def wrapper(*args, **kwargs):
request = kwargs["request"]
handler_task = asyncio.create_task(handler_func(*args, **kwargs))
cancellation_task = asyncio.create_task(listen_for_disconnect(request))
done, pending = await asyncio.wait([handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED)
for task in pending:
task.cancel()
if handler_task in done:
return handler_task.result()
return None
return wrapper
app = FastAPI(lifespan=lifespan)
async def send_request_to_service(
client: httpx.AsyncClient,
prefiller_id: int,
endpoint: str,
req_data: dict,
request_id: str,
max_retries: int = 3,
base_delay: float = 0.2,
):
proxy_state.acquire_aborted_prefiller_requests(prefiller_id)
req_data = req_data.copy()
req_data["stream"] = False
req_data["max_tokens"] = 1
req_data["min_tokens"] = 1
if "max_completion_tokens" in req_data:
req_data["max_completion_tokens"] = 1
if "stream_options" in req_data:
del req_data["stream_options"]
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
last_exc = None
for attempt in range(1, max_retries + 1):
try:
response = await client.post(endpoint, json=req_data, headers=headers)
response.raise_for_status()
if request_id in proxy_state.req_id_future:
result_future = proxy_state.req_id_future[request_id]
result_future.set_result(response.json()["kv_transfer_params"])
return
except (httpx.RequestError, httpx.HTTPStatusError) as e:
logger.warning("Attempt %s failed for %s: %s", attempt, endpoint, e)
last_exc = e
if attempt < max_retries:
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for %s.", max_retries, endpoint)
raise last_exc
async def stream_service_response_with_retry(
client: httpx.AsyncClient,
endpoint: str,
req_data: dict,
request_id: str,
max_retries: int = 3,
base_delay: float = 0.2,
):
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
for attempt in range(1, max_retries + 1):
try:
async with client.stream("POST", endpoint, json=req_data, headers=headers) as response:
response.raise_for_status()
first_chunk_sent = False
async for chunk in response.aiter_bytes():
first_chunk_sent = True
yield chunk
return # Success, exit after streaming
except (httpx.RequestError, httpx.HTTPStatusError) as e:
if attempt < max_retries:
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, e)
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
raise e
except Exception as e:
# If any chunk has been sent, do not retry, just log and drop
if "first_chunk_sent" in locals() and first_chunk_sent:
logger.error("Streaming to client interrupted after response started: %s", e)
return
else:
if attempt < max_retries:
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, e)
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
raise e
def get_api_request_id(api, req_id):
if api == "/completions":
return "cmpl-" + req_id + "-0"
elif api == "/chat/completions":
return "chatcmpl-" + req_id
def get_origin_request_id(api, req_id):
if api == "/completions":
return req_id.replace("cmpl-", "")[:-2]
elif api == "/chat/completions":
return req_id.replace("chatcmpl-", "")
async def _handle_completions(api: str, request: Request):
try:
req_data = await request.json()
req_body = await request.body()
request_length = len(req_body)
request_id = await proxy_state.next_req_id()
request_id_api = get_api_request_id(api, request_id)
proxy_state.req_data_dict[request_id_api] = (copy.deepcopy(req_data), request_length, api)
req_data["kv_transfer_params"] = {
"do_remote_decode": False,
"do_remote_prefill": True,
"metaserver": f"http://{global_args.host}:{global_args.port}/v1/metaserver",
}
# Select decoder
decoder_score = proxy_state.calculate_decode_scores(request_length)
logger.debug("Decoder score: %f", decoder_score)
# Use the prefiller's kv_transfer_params to select decoder
decoder_idx = proxy_state.select_decoder(decoder_score)
decoder = proxy_state.decoders[decoder_idx]
# logger.debug("Using %s %s", prefiller.url, decoder.url)
# Stream response from decoder
released_kv = False
# Record request info for recompute
stream_flag = bool(req_data.get("stream", False))
chat_flag = "messages" in req_data
if "prompt" in req_data:
origin_prompt = req_data["prompt"]
elif chat_flag:
messages = req_data["messages"]
origin_prompt = messages[0].get("content", "")
if isinstance(origin_prompt, list):
origin_prompt = origin_prompt[0].get("text", "")
else:
origin_prompt = ""
# refer to vLLM sampling_params: max_token default value
origin_max_tokens = req_data.get("max_tokens", 16)
async def generate_stream():
nonlocal released_kv
generated_token = ""
released_kv = False
retry_count = 0
retry = True
completion_tokens = 0
# Only one await per chunk, minimal logic in loop
try:
while retry:
retry = False
async for chunk in stream_service_response_with_retry(
decoder.client,
api,
req_data,
request_id=request_id,
max_retries=global_args.max_retries,
base_delay=global_args.retry_delay,
):
try:
chunk_str = chunk.decode("utf-8").strip()
except UnicodeDecodeError:
logger.debug("Skipping chunk: %s", chunk)
yield chunk
continue
if not chunk_str:
continue
if chunk_str.startswith("data: "):
chunk_str = chunk_str[len("data: ") :]
try:
chunk_json = json.loads(chunk_str)
except json.JSONDecodeError:
# if chunk is [done], skip it.
logger.debug("Skipping chunk: %s", chunk_str)
yield chunk
continue
choices = chunk_json.get("choices", [])
if not choices:
yield chunk
continue
choice = choices[0]
delta = choice.get("delta") or {}
message = choice.get("message") or {}
content = delta.get("content") or message.get("content") or choice.get("text") or ""
generated_token += content
stop_reason = choice.get("stop_reason")
usage = chunk_json.get("usage", {})
completion_tokens = (
(completion_tokens + 1)
if stream_flag
else (completion_tokens + usage.get("completion_tokens"))
)
if stop_reason == "recomputed":
retry = True
retry_count += 1
if chat_flag:
messages[0]["content"] = origin_prompt + generated_token
else:
req_data["prompt"] = origin_prompt + generated_token
req_data["max_tokens"] = origin_max_tokens - completion_tokens + retry_count
break
if retry_count > 0 and not stream_flag:
if chat_flag:
choice["message"]["content"] = generated_token
else:
choice["text"] = generated_token
chunk = json.dumps(chunk_json).encode("utf-8")
yield chunk
except Exception as e:
logger.error(
"Error during streaming from decoder %s: %s the aborted request %s "
"will be routing to the target prefiller when new request is ready to dispatch to it",
decoder.url,
e,
request_id,
)
finally:
# After streaming done, release tokens
proxy_state.release_decoder(decoder_idx, decoder_score)
if stream_flag:
return StreamingResponse(generate_stream(), media_type="text/event-stream")
else:
return StreamingResponse(generate_stream(), media_type="application/json")
except Exception as e:
import traceback
exc_info = sys.exc_info()
print(f"Error occurred in disagg prefill proxy server - {api} endpoint")
print(e)
print("".join(traceback.format_exception(*exc_info)))
raise
@app.post("/v1/completions")
@with_cancellation
async def handle_completions(request: Request):
return await _handle_completions("/completions", request)
@app.post("/v1/chat/completions")
@with_cancellation
async def handle_chat_completions(request: Request):
return await _handle_completions("/chat/completions", request)
@app.get("/healthcheck")
async def healthcheck():
return {
"status": "ok",
"prefill_instances": len(proxy_state.prefillers),
"decode_instances": len(proxy_state.decoders),
}
@app.post("/reset_prefix_cache")
async def reset_prefix_cache(request: Request):
params = dict(request.query_params)
failures = []
for client, base_url in [
(s.client, f"http://{s.host}:{s.port}") for s in proxy_state.prefillers + proxy_state.decoders
]:
try:
resp = await client.post(f"{base_url}/reset_prefix_cache", params=params)
resp.raise_for_status()
except Exception as e:
logger.error("reset_prefix_cache failed for %s: %s", base_url, e)
failures.append(base_url)
if failures:
from fastapi.responses import JSONResponse
return JSONResponse(status_code=500, content={"failed": failures})
from fastapi.responses import Response as FastAPIResponse
return FastAPIResponse(status_code=200)
@app.post("/v1/metaserver")
async def metaserver(request: Request):
try:
kv_transfer_params = await request.json()
request_id = kv_transfer_params["request_id"]
assert request_id in proxy_state.req_data_dict
req_data, request_length, api = proxy_state.req_data_dict[request_id]
request_id = get_origin_request_id(api, request_id)
req_data["kv_transfer_params"] = kv_transfer_params
prefiller_score = proxy_state.calculate_prefill_scores(request_length)
logger.debug("Request length: %s, Prefiller score: %s", request_length, prefiller_score)
# Select prefiller
prefiller_idx = proxy_state.select_prefiller(prefiller_score)
prefiller = proxy_state.prefillers[prefiller_idx]
logger.debug("Using prefill prefiller.url=%r req_data=%r", prefiller.url, req_data)
# Send request to prefiller
await send_request_to_service(
prefiller.client,
prefiller_idx,
api,
req_data,
request_id,
max_retries=global_args.max_retries,
base_delay=global_args.retry_delay,
)
except Exception as e:
logger.error("Post metaserver failed with: %s", e)
finally:
proxy_state.release_prefiller(prefiller_idx, prefiller_score)
proxy_state.release_prefiller_kv(prefiller_idx, prefiller_score)
if __name__ == "__main__":
global global_args
global_args = parse_args()
import uvicorn
uvicorn.run(app, host=global_args.host, port=global_args.port)

1293
examples/disaggregated_prefill_v1/load_balance_proxy_server_example.py
View File

File diff suppressed because it is too large Load Diff

45
examples/disaggregated_prefill_v1/mooncake_connector_deployment_guide.md
View File

@@ -1,30 +1,21 @@
# Mooncake connector deployment Guide
## Environmental Dependencies
* Software:
* Python >= 3.9, < 3.12
* CANN >= 8.2.rc1
* PyTorch >= 2.7.1, torch-npu >= 2.7.1.dev20250724
* vLLM (same version as vllm-ascend)
* mooncake-transfer-engine reference documentation: https://github.com/kvcache-ai/Mooncake/blob/main/doc/zh/ascend_transport.md
The vllm version must be the same as the main branch of vllm-ascend, for example, 2025/07/30. The version is
* vllm: v0.10.1
* vllm-ascend: v0.10.1rc1
* vllm: v0.10.1
* vllm-ascend: v0.10.1rc1
## run
### 1.Run `prefill` Node
```
```shell
bash run_prefill.sh
```
Content of the run_prefill.sh script
```
```shell
export HCCL_EXEC_TIMEOUT=204
export HCCL_CONNECT_TIMEOUT=120
export HCCL_IF_IP=localhost
@@ -53,9 +44,7 @@ vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
"kv_role": "kv_producer",
"kv_parallel_size": 1,
"kv_port": "20001",
"engine_id": "0",
"kv_rank": 0,
"kv_connector_module_path": "vllm_ascend.distributed.mooncake_connector",
"kv_connector_extra_config": {
"prefill": {
"dp_size": 2,
@@ -72,28 +61,27 @@ vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
`HCCL_EXEC_TIMEOUT`, `HCCL_CONNECT_TIMEOUT`, and `HCCL_IF_IP` are hccl-related configurations.<br>
Set `GLOO_SOCKET_IFNAME`, `TP_SOCKET_IFNAME`, and `HCCL_SOCKET_IFNAME` to the corresponding NIC.<br>
`ASCEND_RT_VISIBLE_DEVICES` specifies the cards on which the node run resides. The total number of cards equals `dp_size*tp_size`.<br>
`/xxxxx/DeepSeek-V2-Lite-Chat` is configured as a model that requires run.<br>
`/xxxxx/DeepSeek-V2-Lite-Chat` is configured as a model that requires running.<br>
`--host`: indicates the IP address of the node to be started.<br>
`--port`: indicates the port to be started, which corresponds to the port in step 4.<br>
`--seed`, --max-model-len, and --max-num-batched-tokens model basic configuration. Set this parameter based on the site requirements.<br>
`--port`: indicates the port on which the prefill node will listen (e.g., 8100). This port is later referenced in step 3 when configuring the proxy server.<br>
`--seed`: `--max-model-len`, and `--max-num-batched-tokens` are part of the model's basic configuration. Set this parameter based on the site requirements.<br>
`--tensor-parallel-size`: specifies the TP size.<br>
`--data-parallel-size`: indicates the DP size.<br>
`--data-parallel-address`: indicates the IP address of the DP. Set this parameter to the IP address of the node.--data-parallel-rpc-port: indicates the RPC port for communication in the DP group.<br>
`--trust-remote-code` can load the local model.<br>
`--enforce-eager` Turn off the map mode<br>
`--gpu-memory-utilization`: Percentage of video memory occupied by the card<br>
`--kv-transfer-config`: follow kv_connector, kv_connector_module_path: mooncakeconnect, kv_buffer_device, and run on the NPU card. For kv_role, set kv_producer to the p node, kv_consumer to the d node, kv_parallel_size to 1, and kv_port to the port used by the node. For the p node, set engine_id and kv_rank to 0 and for the d node to 1. Configure the distributed parallel policy for the p and d nodes in the kv_connector_extra_config file based on --tensor-parallel-size and --data-parallel-size.<br>
`--kv-transfer-config`: follow kv_connector, kv_buffer_device, and run on the NPU card. For kv_role, set kv_producer to the p node, kv_consumer to the d node, kv_parallel_size to 1, and kv_port to the port used by the node. For the p node, set engine_id and kv_rank to 0 and for the d node to 1. Configure the distributed parallel policy for the p and d nodes in the kv_connector_extra_config file based on --tensor-parallel-size and --data-parallel-size.<br>
### 2. Run `decode` Node
```
```shell
bash run_decode.sh
```
Content of the run_decode.sh script
```
```shell
export HCCL_EXEC_TIMEOUT=204
export HCCL_CONNECT_TIMEOUT=120
export HCCL_IF_IP=localhost
@@ -122,9 +110,7 @@ vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
"kv_role": "kv_consumer",
"kv_parallel_size": 1,
"kv_port": "20002",
"engine_id": "1",
"kv_rank": 1,
"kv_connector_module_path": "vllm_ascend.distributed.mooncake_connector",
"kv_connector_extra_config": {
"prefill": {
"dp_size": 2,
@@ -138,28 +124,27 @@ vllm serve "/xxxxx/DeepSeek-V2-Lite-Chat" \
}'
```
### 3. Start proxy_server. ###
### 3. Start proxy_server
```
```shell
cd /vllm-ascend/examples/disaggregate_prefill_v1/
python load_balance_proxy_server_example.py --host localhost --prefiller-hosts host1 host2 --prefiller-ports 8100 8101 --decoder-hosts host3 host4 --decoder-ports 8200 8201
```
`--host`: indicates the active node. The value of localhost in the curl command delivered in step 5 must be the same as the host. The default port number for starting the service proxy is 8000.<br>
`--prefiller-hosts`: Set this parameter to the IP addresses of all p nodes. In the xpyd scenario, add the IP addresses to the end of this configuration item and leave a blank space between the IP addresses.<br>
`--prefiller-ports`: Set this parameter to the port number of all p nodes, which is the configuration of the port number for the vllm to start the service in step 3. Write the port number after the configuration in sequence and leave a blank space between the port number and the port number. The sequence must be one-to-one mapping to the IP address of --prefiller-hosts.<br>
`--prefiller-ports`: Set this parameter to the port numbers of all prefill (P) nodes, which were defined in step 1 when starting the prefill nodes. Write the port number after the configuration in sequence and leave a blank space between the port number and the port number. The sequence must be one-to-one mapping to the IP address of --prefiller-hosts.<br>
`--decoder-hosts`: Set this parameter to the IP addresses of all d nodes. In the xpyd scenario, add the IP addresses to the end of this configuration item and leave a blank space between the IP addresses.<br>
`--decoder-ports`: Set this parameter to the port number of all d nodes, which is the configuration of the port number for the vllm to start the service in step 4. Set port to the end of the configuration, and leave a blank space between port and port. The sequence must be one-to-one mapping to the IP address of --decoder-hosts.<br>
### 4. Run Inference
Set the IP address in the inference file to the actual IP address. Set the model variable to the path of the model. Ensure that the path is the same as that in the shell script.
```
```shell
curl -s http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{
"model": "model_path",
"prompt": "Given the accelerating impacts of climate change—including rising sea levels, increasing frequency of extreme weather events, loss of biodiversity, and adverse effects on agriculture and human health—there is an urgent need for a robust, globally coordinated response. However, international efforts are complicated by a range of factors: economic disparities between high-income and low-income countries, differing levels of industrialization, varying access to clean energy technologies, and divergent political systems that influence climate policy implementation. In this context, how can global agreements like the Paris Accord be redesigned or strengthened to not only encourage but effectively enforce emission reduction targets? Furthermore, what mechanisms can be introduced to promote fair and transparent technology transfer, provide adequate financial support for climate adaptation in vulnerable regions, and hold nations accountable without exacerbating existing geopolitical tensions or disproportionately burdening those with historically lower emissions?",
"max_tokens": 256
}'
```
```

153
examples/dynamic_bucket_load_balancer/README.md Normal file
View File

@@ -0,0 +1,153 @@
# Dynamic Bucket Load Balancer
A dynamic bucketing-based hybrid load balance proxy for [vLLM](https://github.com/vllm-project/vllm).
The proxy fronts multiple vLLM backend servers and distributes
OpenAI-compatible requests across them. It can run in two modes:
- **Plain load balancing** (default): for each request, estimate a load score
and forward it to the least-loaded backend instance.
- **Dynamic bucket load balancing** (`--enable-dynamic-bucket`): split the
backend pool into a **short-request group** and a **long-request group**, route
requests to a group by their length, and dynamically rebalance across groups
based on the load gap and length affinity.
## Files
- `dynamic_bucket_load_balancer.py` — the core algorithm (pure standard library).
Buckets requests by length, then dynamically adjusts bucket assignment using
bucket load and length affinity.
- `hybrid_proxy_server.py` — the FastAPI proxy server that uses the algorithm to
route requests to the backend servers.
## How It Works
1. **Static bucketing by length.** Each request is first mapped to its *standard
bucket* by request length. With dynamic bucketing enabled the proxy uses two
buckets: short `[0, --server-group-threshold)` and long
`[--server-group-threshold, --max-request-tokens)`.
2. **Server groups.** The ordered backend list is split into the same number of
groups as buckets, **in order**: the first instances form the short group, the
last instances form the long group. With 4 backends and 2 buckets, backends 0
and 1 serve the short bucket, backends 2 and 3 serve the long bucket.
- > **Tip:** configure the first two instances for short sequences and the
> last two for long sequences (e.g. smaller `max-model-len` / KV cache for
> the short group, larger for the long group) to get the best throughput.
3. **Dynamic rebalancing.** For a new request, the balancer looks at neighbor
buckets with a lighter load and computes a redirect probability
`(load-gap probability) × (length-affinity factor)`. If it exceeds the
threshold (`0.12`), the request is redirected to the neighbor bucket. This
means a large load gap is suppressed when the request length is far from the
neighbor bucket, while a modest gap can still trigger a redirect when the
length is close to the boundary.
4. **Within a group**, the least-loaded server (smallest active token count) is
picked via a min-heap, the load is accumulated for the duration of the
request, and released when streaming completes.
## Prerequisites
- Python 3.10+
- Install dependencies:
```bash
pip install "fastapi<0.124.0" httpx uvicorn
```
## Step 1: Start Your Backend Servers
Start at least two vLLM servers, each as a separate process on its own port. The
proxy also works with a single backend, but load balancing is only meaningful
with two or more.
```bash
vllm serve --host 0.0.0.0 --port 8100 ... # vLLM Server 0
vllm serve --host 0.0.0.0 --port 8101 ... # vLLM Server 1
```
## Step 2: Start the Proxy Server
From `examples/dynamic_bucket_load_balancer/`, point the proxy at each backend
with `--server-hosts` / `--server-ports`:
```bash
python hybrid_proxy_server.py \
--host 0.0.0.0 --port 8000 \
--server-hosts 127.0.0.1 127.0.0.1 \
--server-ports 8100 8101
```
This starts the proxy on port 8000 and load balances across the two backends.
### Enable Dynamic Bucket Load Balancing
Add `--enable-dynamic-bucket` to split the pool into short/long groups. The
server count must be `>= 2` so each bucket has at least one instance. With 4
servers the first two form the short group and the last two the long group:
```bash
python hybrid_proxy_server.py \
--host 0.0.0.0 --port 8000 \
--server-hosts 127.0.0.1 127.0.0.1 127.0.0.1 127.0.0.1 \
--server-ports 8100 8101 8102 8103 \
--enable-dynamic-bucket \
--server-group-threshold 32768
```
## Step 3: Send a Request to the Proxy
Send OpenAI-compatible requests to the proxy. For example:
```bash
curl -X POST http://localhost:8000/v1/completions \
-H "Content-Type: application/json" \
-d '{
"model": "your-model",
"prompt": "The quick brown fox jumps over the lazy dog",
"max_tokens": 16
}'
```
Or for chat completions:
```bash
curl -X POST http://localhost:8000/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "your-model",
"messages": [{"role": "user", "content": "Hello!"}],
"max_tokens": 16
}'
```
## Step 4: Health Check
Check that the proxy is running and how many backends it fronts:
```bash
curl http://localhost:8000/healthcheck
```
Returns a JSON object, e.g.:
```json
{"status": "ok", "server_instances": 2}
```
## Configuration
| Argument | Default | Description |
| --- | --- | --- |
| `--host` | `localhost` | Proxy listen host. |
| `--port` | `8000` | Proxy listen port. |
| `--server-hosts` | `localhost` | Hosts of the backend vLLM servers (one per server, in order). |
| `--server-ports` | `8001` | Ports of the backend vLLM servers (one per server, in order). |
| `--enable-dynamic-bucket` | `False` | Enable dynamic bucket load balancing. |
| `--server-group-threshold` | `32768` | Length boundary between the short and long buckets. |
| `--max-request-tokens` | `131072` | Upper bound of the long bucket (max request length). |
| `--max-retries` | `3` | Max retries for a backend HTTP request. |
| `--retry-delay` | `0.001` | Base delay (seconds) for exponential backoff retries. |
The number of `--server-hosts` must equal the number of `--server-ports`.

243
examples/dynamic_bucket_load_balancer/dynamic_bucket_load_balancer.py Normal file
View File

@@ -0,0 +1,243 @@
import math
from collections import namedtuple
from typing import AnyStr
ServerInfo = namedtuple("ServerInfo", ["instance_type", "instance_idx"])
class Task:
"""Request task (carries request length info)."""
def __init__(self, task_id, task_length, task_load):
self.id = task_id
self.length = task_length
self.bucket_idx = -1
self.load = task_load
self.server_info: ServerInfo = ServerInfo("Unknown", -1)
def __repr__(self):
return (
f"Task(id={self.id}, length={self.length}, load={self.load}, "
f"instance_type={self.server_info.instance_type}, instance_idx={self.server_info.instance_idx})"
)
class Bucket:
"""A bucket."""
def __init__(self, bucket_ranges: tuple[int, int]):
self.min_length = bucket_ranges[0]
self.max_length = bucket_ranges[1]
self.task_count = 0
self.total_load = 0.0
class DynamicBucketLoadBalancer:
"""
Statically buckets requests by length first, then dynamically adjusts the
assignment of new requests based on bucket load and length affinity to
achieve load balancing.
"""
def __init__(
self, buckets: list[tuple[int, int]], sensitivity=1.0, affinity_strength=0.1, log_func=print, all_neighbor=False
):
"""
Initialize the load balancer.
:param buckets: length range of each bucket
:param sensitivity: sensitivity to the load gap (higher = more sensitive)
:param affinity_strength: strength of length affinity (higher = a request
stays more in its standard bucket)
:param log_func: logging function
:param all_neighbor: if False, only the left/right buckets are neighbors;
if True, all buckets are neighbors
"""
self.num_buckets = len(buckets)
self.sensitivity = sensitivity
self.affinity_strength = affinity_strength
self.log_func = log_func
self.all_neighbor = all_neighbor
self.buckets = {idx: Bucket(bucket_ranges) for idx, bucket_ranges in enumerate(buckets)}
bucket_boundaries = ", ".join(
f"bucket {idx}: [{bucket.min_length}, {bucket.max_length})" for idx, bucket in self.buckets.items()
)
self._log_info(f"Initialized {self.num_buckets} buckets: {bucket_boundaries}")
# Redirect only when the redirect probability exceeds this threshold
self.base_probability_threshold = 0.12
self._log_info(f"Load Balance base_probability_threshold: {self.base_probability_threshold:.2f} ")
# Track request tasks
self.tasks: dict[AnyStr, Task] = {} # type: ignore
# Statistics
self.redirected_tasks = 0
self.total_tasks = 0
def _log_info(self, msg, *args, **kwargs):
if self.log_func:
self.log_func(msg, *args, **kwargs)
def _get_standard_bucket_index(self, task_length):
"""Return the standard bucket index for the given request length."""
for bucket_idx, bucket in self.buckets.items():
if bucket.min_length <= task_length < bucket.max_length:
return bucket_idx
# Fall back to the last bucket if the length is outside every range
return self.num_buckets - 1
def _get_neighbor_indices(self, bucket_idx):
"""Return the left and right neighbor indices of the given bucket."""
if self.all_neighbor:
return list(range(self.num_buckets))
neighbors = []
if bucket_idx > 0:
neighbors.append(bucket_idx - 1)
if bucket_idx < self.num_buckets - 1:
neighbors.append(bucket_idx + 1)
return neighbors
def _calculate_length_affinity(self, task_length, neighbor_bucket_idx):
"""
Compute the affinity factor between the task length and the neighbor
bucket (0.0 to 1.0). 1.0 means right next to the neighbor bucket, 0.0
means far away from it.
"""
neighbor_bucket = self.buckets[neighbor_bucket_idx]
neighbor_bucket_min = neighbor_bucket.min_length
neighbor_bucket_max = neighbor_bucket.max_length
if neighbor_bucket_min < task_length < neighbor_bucket_max:
raise RuntimeError("task_length must be outside the neighbor bucket range")
neighbor_bucket_center = (neighbor_bucket_min + neighbor_bucket_max) / 2.0
neighbor_bucket_half_width = (neighbor_bucket_max - neighbor_bucket_min) / 2.0
distance_to_center = abs(task_length - neighbor_bucket_center)
# The closer to the neighbor bucket boundary, the closer the affinity to 1
if neighbor_bucket_half_width > 0:
# Relative distance from the neighbor bucket half-width
normalized_distance = (distance_to_center - neighbor_bucket_half_width) / neighbor_bucket_half_width
# Exponential decay, e.g. normalized_distance=0.1, affinity_strength=1.0 -> 0.9
neighbor_affinity = math.exp(-self.affinity_strength * normalized_distance)
else:
neighbor_affinity = 1.0 # Safeguard; unreachable in practice
# Clamp to [0, 1]
return max(0.0, min(neighbor_affinity, 1.0))
def _calculate_redirect_probability(self, task_length, standard_bucket_idx, neighbor_bucket_idx):
"""
Compute the probability of redirecting to the neighbor bucket based on
the load gap and length affinity.
"""
standard_load = self.buckets[standard_bucket_idx].total_load
neighbor_load = self.buckets[neighbor_bucket_idx].total_load
# --- 1. Base probability from the load gap ---
if standard_load <= 0:
load_probability = 0.0 # No load in the standard bucket -> no redirect
else:
load_ratio = neighbor_load / max(standard_load, 1e-9) # Guard against division by zero
# The smaller the neighbor load relative to the standard load, the
# higher the redirect probability
# e.g. neighbor/standard = 5/6, sensitivity=1.0 -> 1/6 ≈ 0.1667
load_probability = 1 - load_ratio**self.sensitivity
load_probability = max(0.0, min(load_probability, 1))
# --- 2. Length affinity factor ---
affinity_factor = self._calculate_length_affinity(task_length, neighbor_bucket_idx)
# --- 3. Final probability: load gap * length affinity ---
# A large load gap is suppressed when the request length is far from the
# neighbor bucket; a modest gap can still redirect when it is close.
final_probability = load_probability * affinity_factor
return final_probability
def dispatch_single_task(self, task_id: AnyStr, task_length: int, task_load):
return self.dispatch_task(Task(task_id, task_length, task_load))
def dispatch_task(self, cur_task):
"""
Assign a bucket to a new request, considering dynamic load balancing and
length affinity.
"""
self.total_tasks += 1
standard_bucket_idx = self._get_standard_bucket_index(cur_task.length)
neighbor_indices = self._get_neighbor_indices(standard_bucket_idx)
best_neighbor_idx = None
best_redirect_prob = 0.0
# Pick the neighbor with the highest redirect probability
for neighbor_idx in neighbor_indices:
# Only consider neighbors with lower load
if self.buckets[neighbor_idx].total_load < self.buckets[standard_bucket_idx].total_load:
prob = self._calculate_redirect_probability(cur_task.length, standard_bucket_idx, neighbor_idx)
if prob > best_redirect_prob:
best_redirect_prob = prob
best_neighbor_idx = neighbor_idx
# Decide the final bucket
final_bucket_idx = standard_bucket_idx
if best_neighbor_idx is not None and best_redirect_prob > 0:
if self.base_probability_threshold < best_redirect_prob:
final_bucket_idx = best_neighbor_idx
self.redirected_tasks += 1
self._log_info(
f"{cur_task} redirected from bucket {standard_bucket_idx} to {final_bucket_idx}"
f"(prob={best_redirect_prob:.4f})"
)
# Bookkeeping on the chosen bucket
self.buckets[final_bucket_idx].task_count += 1
self.buckets[final_bucket_idx].total_load += cur_task.load
cur_task.bucket_idx = final_bucket_idx
if cur_task.id in self.tasks:
raise RuntimeError(f"Task {cur_task.id} is existed!")
else:
self.tasks[cur_task.id] = cur_task
return final_bucket_idx, cur_task
def release_task(self, task_id):
"""Release the load of a request."""
if task_id in self.tasks:
found_task = self.tasks.pop(task_id)
if 0 <= found_task.bucket_idx < self.num_buckets:
self.buckets[found_task.bucket_idx].task_count -= 1
self.buckets[found_task.bucket_idx].total_load -= found_task.load
return True
else:
raise RuntimeError(f"Bucket {found_task.bucket_idx} not found")
else:
raise RuntimeError(f"Task {task_id} not found")
def release_all_tasks(self):
for bucket in self.buckets.values():
bucket.task_count = 0
bucket.total_load = 0
self.tasks.clear()
class NoStandardBucketLoadBalancer(DynamicBucketLoadBalancer):
"""Dispatch requests by load only, with no standard bucket."""
def __init__(self, num_buckets: int, max_length: int, log_func=print):
bucket_range = math.ceil(max_length / num_buckets)
start_length = 0
buckets = []
for _ in range(num_buckets):
end_length = start_length + bucket_range
if end_length > max_length:
end_length = max_length
buckets.append((start_length, end_length))
start_length += bucket_range
super().__init__(buckets=buckets, log_func=log_func, sensitivity=100, affinity_strength=0, all_neighbor=True)

455
examples/dynamic_bucket_load_balancer/hybrid_proxy_server.py Normal file
View File

@@ -0,0 +1,455 @@
# Adapted from https://github.com/vllm-project/vllm/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
# SPDX-License-Identifier: Apache-2.0
#
# Dynamic bucketing-based hybrid load balance proxy server.
# See README.md in this directory for the tutorial and usage.
import argparse
import asyncio
import functools
import heapq
import os
import sys
import uuid
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any
import httpx
from dynamic_bucket_load_balancer import DynamicBucketLoadBalancer, ServerInfo, Task
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
try:
from vllm.logger import init_logger
logger = init_logger(__name__)
except ImportError:
import logging
logger = logging.getLogger(__name__)
# Use uvloop for a faster event loop if available
try:
import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
except ImportError:
pass
class ServerState:
def __init__(self, host, port):
self.host = host
self.port = port
self.url = f"http://{host}:{port}/v1"
self.client = httpx.AsyncClient(
timeout=None,
base_url=self.url,
limits=httpx.Limits(max_connections=100000, max_keepalive_connections=100000),
)
self.active_tokens = 0
self.aborted_requests = set()
def __eq__(self, other):
self_host = self.host.replace("localhost", "0.0.0.0").replace("127.0.0.1", "0.0.0.0")
other_host = other.host.replace("localhost", "0.0.0.0").replace("127.0.0.1", "0.0.0.0")
return self_host == other_host and str(self.port) == str(other.port)
def __hash__(self):
self_host = self.host.replace("localhost", "0.0.0.0").replace("127.0.0.1", "0.0.0.0")
return hash((self_host, str(self.port)))
def __repr__(self):
return f"{self.host}:{self.port}"
@dataclass(order=True)
class ServerHeapItem:
priority: float
server_idx: int
server: ServerState
class ProxyState:
def __init__(self, server_instances):
self.infer_servers: list[ServerState] = [ServerState(h, p) for h, p in server_instances]
self.req_id_lock = asyncio.Lock()
# Dynamic bucket load balancer
self.bucket_load_balancer = None
if global_args.enable_dynamic_bucket:
self.num_buckets = 2 # Two buckets (short/long) when dynamic bucketing is enabled
self.server_group_threshold = global_args.server_group_threshold
buckets = [(0, self.server_group_threshold), (self.server_group_threshold, global_args.max_request_tokens)]
self.bucket_load_balancer = DynamicBucketLoadBalancer(buckets=buckets)
else:
self.num_buckets = 1 # No bucketing by default
# Priority queue per group; smaller score = higher priority (lower load)
server_heap_items = [ServerHeapItem(0.0, i, server) for i, server in enumerate(self.infer_servers)]
self.server_heaps: list[list[ServerHeapItem]] = self._group_servers(server_heap_items, self.num_buckets)
self.server_idx_to_group_idx = {}
# Heapify each group
for idx, cur_heap in enumerate(self.server_heaps):
for server_item in cur_heap:
self.server_idx_to_group_idx[server_item.server_idx] = idx
heapq.heapify(cur_heap)
logger.info(
"Dynamic bucket enabled: %s, number of groups: %s",
global_args.enable_dynamic_bucket,
len(self.server_heaps),
)
for group_idx, cur_heap in enumerate(self.server_heaps):
logger.info("Group %s: %s", group_idx, cur_heap)
@staticmethod
def _group_servers(servers: list[ServerHeapItem], num_groups: int):
"""
Split servers into num_groups groups.
Args:
servers (list): the server list to group.
num_groups (int): the number of groups.
Returns:
list[list]: the grouped server list.
Raises:
ValueError: when num_groups <= 0.
"""
if num_groups <= 0:
raise ValueError("Num of group is illegal")
if len(servers) < num_groups:
raise ValueError("Number of servers must greater than or equal to number of groups")
n = len(servers)
if n == 0:
return [[] for _ in range(num_groups)]
elif n == 1:
return [servers]
base_size = n // num_groups
remainder = n % num_groups
groups = []
start_index = 0
for i in range(num_groups):
group_size = base_size + 1 if i < remainder else base_size
end_index = start_index + group_size
groups.append(servers[start_index:end_index])
start_index = end_index
return groups
def _update_server_priority(self, server_idx: int):
"""Update the priority of a server in the heap."""
server = self.infer_servers[server_idx]
priority = server.active_tokens
# Remove the old entry, then add the new one
group_idx = self.server_idx_to_group_idx[server_idx]
self.server_heaps[group_idx] = [
server_heap_item
for server_heap_item in self.server_heaps[group_idx]
if server_heap_item.server_idx != server_idx
]
self.server_heaps[group_idx].append(ServerHeapItem(priority, server_idx, server))
heapq.heapify(self.server_heaps[group_idx])
async def next_req_id(self):
async with self.req_id_lock:
return str(uuid.uuid4())
def select_server(self, token_count, group_idx: int):
if not self.infer_servers:
raise RuntimeError("No inference servers available")
server_heap_item: ServerHeapItem = heapq.heappop(self.server_heaps[group_idx])
chosen_server_idx = server_heap_item.server_idx
# Update the chosen server (accumulate load)
self.infer_servers[chosen_server_idx].active_tokens += token_count
# Update priority and re-add to the heap
self._update_server_priority(chosen_server_idx)
return chosen_server_idx
def release_server(self, idx: int, token_count, req_id):
self.infer_servers[idx].active_tokens -= token_count
if global_args.enable_dynamic_bucket and req_id is not None and self.bucket_load_balancer is not None:
self.bucket_load_balancer.release_task(req_id)
# Update the priority queue after release
self._update_server_priority(idx)
def calculate_request_score(self, request_length: int, max_tokens: int = 16, ignore_eos: bool = False) -> float:
if ignore_eos:
return request_length + max_tokens
else:
# Note that 0.5 is an empirical value here because we don't know
# the actual number of tokens generated before EOS.
return request_length + 0.5 * max_tokens
def calculate_request_tokens(self, request_length: int) -> float:
return request_length / 4.0
def select_server_group(self, req_id: str, request_tokens, priority_score) -> tuple[int, Task | None]:
"""Pick the best group given the request length and the current load of each group."""
if global_args.enable_dynamic_bucket and self.bucket_load_balancer is not None:
group_idx, task = self.bucket_load_balancer.dispatch_single_task(req_id, request_tokens, priority_score)
return group_idx, task
else:
return 0, None
proxy_state = None
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--host", type=str, default="localhost")
parser.add_argument("--server-hosts", type=str, nargs="+", default=["localhost"])
parser.add_argument("--server-ports", type=int, nargs="+", default=[8001])
parser.add_argument("--max-retries", type=int, default=3, help="Maximum number of retries for HTTP requests")
parser.add_argument(
"--retry-delay", type=float, default=0.001, help="Base delay (seconds) for exponential backoff retries"
)
parser.add_argument("--server-group-threshold", type=int, default=32 * 1024, help="Threshold of server groups")
parser.add_argument("--max-request-tokens", type=int, default=128 * 1024, help="Max tokens of request")
parser.add_argument(
"--enable-dynamic-bucket", action="store_true", default=False, help="Enable dynamic bucket load Balancer"
)
args = parser.parse_args()
if len(args.server_hosts) != len(args.server_ports):
raise ValueError("Number of dp hosts must match number of dp ports")
args.server_instances = list(zip(args.server_hosts, args.server_ports))
return args
@asynccontextmanager
async def lifespan(app: FastAPI):
global proxy_state
proxy_state = ProxyState(global_args.server_instances)
logger.debug("Initialized %s dp server clients.", len(proxy_state.infer_servers))
yield
for p in proxy_state.infer_servers:
await p.client.aclose()
async def listen_for_disconnect(request: Request) -> None:
"""Return when a disconnect message is received."""
while True:
message = await request.receive()
if message["type"] == "http.disconnect":
break
def with_cancellation(handler_func):
@functools.wraps(handler_func)
async def wrapper(*args, **kwargs):
request = kwargs["request"]
handler_task = asyncio.create_task(handler_func(*args, **kwargs))
cancellation_task = asyncio.create_task(listen_for_disconnect(request))
done, pending = await asyncio.wait([handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED)
for task in pending:
task.cancel()
if handler_task in done:
return handler_task.result()
return None
return wrapper
app = FastAPI(lifespan=lifespan)
async def stream_service_response_with_retry(
client: httpx.AsyncClient,
endpoint: str,
req_data: dict,
request_id: str,
max_retries: int = 3,
base_delay: float = 0.2,
):
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
for attempt in range(1, max_retries + 1):
# Reset per retry to avoid leaking a stale True from a previous iteration
first_chunk_sent = False
try:
async with client.stream("POST", endpoint, json=req_data, headers=headers) as response:
response.raise_for_status()
async for chunk in response.aiter_bytes():
first_chunk_sent = True
yield chunk
return # Success; exit after streaming completes
except (httpx.RequestError, httpx.HTTPStatusError) as e:
# After the first chunk is forwarded, retry is forbidden (would duplicate/corrupt the stream).
if first_chunk_sent:
logger.error("Streaming to client interrupted after response started: %s", str(e))
return
if attempt < max_retries:
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, str(e))
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
raise e
except Exception as e:
# Same guard as above for non-HTTP exceptions
if first_chunk_sent:
logger.error("Streaming to client interrupted after response started: %s", str(e))
return
if attempt < max_retries:
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, str(e))
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
raise e
async def _select_instance(api: str, req_data: Any, request_length: int):
# refer to vLLM sampling_params: max_token default value
max_tokens = req_data.get("max_tokens", 16)
ignore_eos = req_data.get("ignore_eos", False)
priority_score = 0.0
if global_args.enable_dynamic_bucket:
priority_score = proxy_state.calculate_request_tokens(request_length)
else:
priority_score = proxy_state.calculate_request_score(
request_length, max_tokens=max_tokens, ignore_eos=ignore_eos
)
logger.debug(
"Request length: %s, max tokens: %s, ignore_eos: %s, Priority score: %s",
request_length,
max_tokens,
ignore_eos,
priority_score,
)
request_id = await proxy_state.next_req_id()
# Select server based on priority score
request_tokens = proxy_state.calculate_request_tokens(request_length)
group_idx, task = proxy_state.select_server_group(request_id, request_tokens, priority_score)
try:
server_idx = proxy_state.select_server(priority_score, group_idx)
except Exception:
if global_args.enable_dynamic_bucket and task is not None and proxy_state.bucket_load_balancer is not None:
proxy_state.bucket_load_balancer.release_task(task.id)
raise
if global_args.enable_dynamic_bucket and task is not None:
task.server_info = ServerInfo("DP", server_idx)
chosen_server = proxy_state.infer_servers[server_idx]
logger.debug(
"[group_idx=%s, server_idx=%s] Choose server %s to process request %s",
group_idx,
server_idx,
chosen_server.url,
request_id,
)
return InstanceInfo(
request_id=request_id, server_idx=server_idx, priority_score=priority_score, server_state=chosen_server
)
@dataclass
class InstanceInfo:
request_id: str
server_idx: int
priority_score: float
server_state: ServerState
async def _handle_completions(api: str, request: Request):
# streaming_started ensures release_server runs exactly once: in
# generate_stream's finally on the normal path, or below if it never started.
instance_info = None
streaming_started = False
try:
req_data = await request.json()
req_body = await request.body()
request_length = len(req_body)
instance_info = await _select_instance(api, req_data, request_length)
async def generate_stream():
nonlocal instance_info
try:
async for chunk in stream_service_response_with_retry(
instance_info.server_state.client, # type: ignore
api,
req_data,
request_id=instance_info.request_id, # type: ignore
max_retries=global_args.max_retries,
base_delay=global_args.retry_delay,
):
yield chunk
except Exception as e:
logger.error(
"Error during streaming from server %s: %s, the aborted request is: %s.",
instance_info.server_state.url, # type: ignore
str(e),
instance_info.request_id, # type: ignore
)
finally:
# Release load after streaming completes
proxy_state.release_server( # type: ignore
instance_info.server_idx, # type: ignore
instance_info.priority_score, # type: ignore
instance_info.request_id, # type: ignore
)
streaming_started = True
return StreamingResponse(generate_stream(), media_type="application/json")
except Exception as e:
import traceback
exc_info = sys.exc_info()
print(f"Error occurred in external dp proxy server - {api} endpoint")
print(e)
print("".join(traceback.format_exception(*exc_info)))
raise
finally:
# If streaming never started (client disconnect or selection error),
# release here to avoid leaking active_tokens / the bucket task; the
# normal path already released in generate_stream.
if instance_info is not None and not streaming_started:
proxy_state.release_server(instance_info.server_idx, instance_info.priority_score, instance_info.request_id)
@app.post("/v1/completions")
@with_cancellation
async def handle_completions(request: Request):
return await _handle_completions("/completions", request)
@app.post("/v1/chat/completions")
@with_cancellation
async def handle_chat_completions(request: Request):
return await _handle_completions("/chat/completions", request)
@app.get("/healthcheck")
async def healthcheck():
return {
"status": "ok",
"server_instances": len(proxy_state.infer_servers),
}
if __name__ == "__main__":
global global_args
global_args = parse_args()
import uvicorn
uvicorn.run(app, host=global_args.host, port=global_args.port)

304
examples/epd_disaggregated/epd_disaggregated_guide.md Normal file
View File

@@ -0,0 +1,304 @@
# EPD disaggregated deployment Guide
## run
The EPD disaggregated technology accelerates model inference by decoupling the visual encoding computation and LLM computation stages. Currently, the EPD separation feature can achieve different data transmissions between E and P/PD nodes by configuring different connector backends. Vllm-ascend currently supports the ECExampleConnector backend implemented on vllm, and will support Mooncake as well as shared memory(SHM) backend transmission methods in the future.
### ECexample-connector deployment guide
Using the Qwen3-VL-8B model inference as an example.
#### 1. run 1e1pd case
##### 1.1 run e node
```shell
bash run_e.sh
```
Content of the run_e.sh script
```shell
unset ftp_proxy
unset https_proxy
unset http_proxy
EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/data/ec_cache}"
rm /data/ec_cache -rf
mkdir -p /data/ec_cache
export ASCEND_RT_VISIBLE_DEVICES=0
vllm serve "/your/local/model/path/Qwen3-VL-8B-Instruct" \
--gpu-memory-utilization 0.01 \
--port "23001" \
--enforce-eager \
--enable-request-id-headers \
--served-model-name qwenvl \
--max-model-len 32768 \
--no-enable-prefix-caching \
--max-num-batched-tokens 114688 \
--max-num-seqs 128 \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_producer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}'
```
`--gpu-memory-utilization`:For LLM Model, It is usually used to control the kv cache allocation.For model architectures like vision encoder that do not require KV Cache, it is usually set to 0.01 to minimize HBM usage.<br>
`--ec-transfer-config`:Specify ec-transfer connector settings.For ECExampleConnector, you need to specify the role played by the current node(For e node, set it to 'ec_producer') and the local memory address for data transfer between nodes.<br>
##### 1.2 run pd node
```shell
bash run_pd.sh
```
Content of the run_pd.sh script
```shell
unset ftp_proxy
unset https_proxy
unset http_proxy
EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/data/ec_cache}"
export ASCEND_RT_VISIBLE_DEVICES=1
vllm serve "/your/local/model/path/Qwen3-VL-8B-Instruct" \
--gpu-memory-utilization 0.7 \
--port "33005" \
--enforce-eager \
--enable-request-id-headers \
--served-model-name qwenvl \
--max-model-len 32768 \
--max-num-seqs 128 \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_consumer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}'
```
`--ec-transfer-config`:Same as e node,but ec_role is set to 'ec_consumer'.<br>
##### 1.3 run proxy node
```shell
bash run_proxy.sh
```
Content of the run_proxy.sh script
```bash
python3 epd_load_balance_proxy_layerwise_server_example.py \
--encoder-hosts 127.0.0.1 \
--encoder-ports 23001 \
--pd-hosts 127.0.0.1 \
--pd-ports 33005 \
--host 127.0.0.1 \
--port 8001
```
The parameters are explained as follows:<br>
`--encoder-hosts`: E node IP address.<br>
`--encoder-ports`: The E node port number. It needs to be consistent with the --port in the E node's startup script.<br>
`--pd-hosts`: PD node IP address.<br>
`--pd-ports`: The PD node port number. It needs to be consistent with the --port in the PD node's startup script.<br>
`--host`: Proxy node IP address.<br>
`--port`: Proxy node port number.<br>
##### 1.4 run inference
```bash
curl http://localhost:8001/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "qwenvl",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": "https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png"}},
{"type": "text", "text": "What is the text in the illustration?"}
]}
]
}'
```
#### 2.run 1e1p1d case
##### 2.1 run e node
```shell
bash run_e.sh
```
Content of the run_e.sh script
```shell
unset ftp_proxy
unset https_proxy
unset http_proxy
EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/data/ec_cache}"
rm /data/ec_cache -rf
mkdir -p /data/ec_cache
export ASCEND_RT_VISIBLE_DEVICES=0
vllm serve "/home/p00929506/Qwen3-VL-8B-Instruct" \
--gpu-memory-utilization 0.01 \
--port "23001" \
--enforce-eager \
--enable-request-id-headers \
--served-model-name qwenvl \
--max-model-len 32768 \
--no-enable-prefix-caching \
--max-num-batched-tokens 114688 \
--max-num-seqs 128 \
--ec-transfer-config '{
"ec_connector": "ECExampleConnector",
"ec_role": "ec_producer",
"ec_connector_extra_config": {
"shared_storage_path": "'"$EC_SHARED_STORAGE_PATH"'"
}
}'
```
##### 2.2 run p node
```shell
bash run_p.sh
```
Content of the run_p.sh script
```shell
unset ftp_proxy
unset https_proxy
unset http_proxy
EC_SHARED_STORAGE_PATH="${EC_SHARED_STORAGE_PATH:-/data/ec_cache}"
export ASCEND_RT_VISIBLE_DEVICES=1
vllm serve "/home/p00929506/Qwen3-VL-8B-Instruct" \
--gpu-memory-utilization 0.7 \
--port "33003" \
--enforce-eager \
--enable-request-id-headers \
--served-model-name qwenvl \
--max-model-len 32768 \
--max-num-seqs 128 \
--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": "MooncakeLayerwiseConnector",
"kv_role": "kv_producer",
"kv_port": "50001",
"kv_connector_extra_config": {
"use_ascend_direct": true,
"prefill": {
"dp_size": 1,
"tp_size": 1
},
"decode": {
"dp_size": 1,
"tp_size": 1
}
}
}'
```
##### 2.3 run d node
```shell
bash run_d.sh
```
Content of the run_d.sh script
```shell
unset ftp_proxy
unset https_proxy
unset http_proxy
export ASCEND_RT_VISIBLE_DEVICES=4
vllm serve "/your/local/model/path/Qwen3-VL-8B-Instruct" \
--gpu-memory-utilization 0.7 \
--port "33006" \
--enforce-eager \
--enable-request-id-headers \
--served-model-name qwenvl \
--max-model-len 32768 \
--max-num-seqs 128 \
--kv-transfer-config \
'{"kv_connector": "MooncakeLayerwiseConnector",
"kv_role": "kv_consumer",
"kv_port": "50001",
"kv_connector_extra_config": {
"use_ascend_direct": true,
"prefill": {
"dp_size": 1,
"tp_size": 1
},
"decode": {
"dp_size": 1,
"tp_size": 1
}
}
}'
```
##### 2.4 run proxy node
```shell
bash run_proxy.sh
```
Content of the run_proxy.sh script
```shell
python3 epd_load_balance_proxy_layerwise_server_example.py \
--encoder-hosts 127.0.0.1 \
--encoder-ports 23001 \
--prefiller-hosts 127.0.0.1 \
--prefiller-ports 33003 \
--decoder-hosts 127.0.0.1 \
--decoder-ports 33006 \
--host 127.0.0.1 \
--port 8001
```
`--prefiller-hosts`: Prefill node IP address.<br>
`--prefiller-ports`: The Prefill node port number. It needs to be consistent with the --port in the Prefill node's startup script.<br>
`--decoder-hosts`: Decode node IP address.<br>
`--decoder-ports`: The Decode node port number. It needs to be consistent with the --port in the Decode node's startup script.<br>
##### 2.5 run inference
```bash
curl http://localhost:8001/v1/chat/completions \
-H "Content-Type: application/json" \
-d '{
"model": "qwenvl",
"messages": [
{"role": "system", "content": "You are a helpful assistant."},
{"role": "user", "content": [
{"type": "image_url", "image_url": {"url": "https://modelscope.oss-cn-beijing.aliyuncs.com/resource/qwen.png"}},
{"type": "text", "text": "What is the text in the illustration?"}
]}
]
}'
```

911
examples/epd_disaggregated/epd_load_balance_proxy_layerwise_server_example.py Normal file
View File

@@ -0,0 +1,911 @@
# Adapted from https://github.com/vllm-project/vllm-ascend/blob/main/examples/disaggregated_prefill_v1/load_balance_proxy_layerwise_server_example.py
# SPDX-License-Identifier: Apache-2.0
#
# Tutorial: Using the EPD Load Balance Proxy Server Example
#
# This proxy server is designed to distribute requests between multiple
# "encoder", "pd", "prefiller" and "decoder" backend servers for large language model inference.
# It is useful for scaling out inference workloads and balancing load across
# multiple backend instances.
#
# Features:
# - Load balances multimodal requests to multiple encoder, pd, prefiller and decoder servers.
# - Supports OpenAI-compatible /v1/completions and /v1/chat/completions endpoints.
# - Streams responses from backend servers to clients.
#
# Prerequisites:
# - Python 3.8+
# - Install dependencies:
# pip install fastapi<0.124.0 httpx uvicorn vllm
#
# Step 1: Start Your Backend Servers
# ----------------------------------
# You need to have at least one prefiller and one decoder backend running.
# These can be mock servers or actual vLLM servers.
#
# For testing, you can use the provided mock server:
#
# vllm serve --host 0.0.0.0 --port 8101 ... # Encoder 1
# vllm serve --host 0.0.0.0 --port 8102 ... # Encoder 2
# vllm serve --host 0.0.0.0 --port 8201 ... # PD 1
# vllm serve --host 0.0.0.0 --port 8202 ... # PD 2
# vllm serve --host 0.0.0.0 --port 8301 ... # Prefiller 1
# vllm serve --host 0.0.0.0 --port 8301 ... # Prefiller 2
# vllm serve --host 0.0.0.0 --port 8401 ... # Decoder 1
# vllm serve --host 0.0.0.0 --port 8402 ... # Decoder 2
#
# Step 2: Start the Proxy Server
# ------------------------------
# Run the proxy server, specifying the host/port for each instance:
#
# 2 Encoder instance + 2 PD instance:
# python epd_load_balance_proxy_layerwise_server_example.py \
# --encoder-hosts 127.0.0.1 127.0.0.1 \
# --encoder-ports 81001 81002 \
# --pd-hosts 127.0.0.1 127.0.0.1 \
# --pd-ports 82001 82002 \
# --host 0.0.0.0 \
# --port 9000
# 2 Encoder instance + 2 Prefill instance + 2 Decode instance:
# python epd_load_balance_proxy_layerwise_server_example.py \
# --encoder-hosts 127.0.0.1 127.0.0.1 \
# --encoder-ports 81001 81002 \
# --prefiller-hosts 127.0.0.1 127.0.0.1 \
# --prefiller-ports 83001 83002 \
# --decoder-hosts 127.0.0.1 127.0.0.1 \
# --decoder-ports 84001 84002 \
# --host 0.0.0.0 \
# --port 9000
# This will start the proxy on port 9000, load balancing between two encoder, tweo pd, two prefiller
# and two decoder servers.
#
# Step 3: Send a Request to the Proxy
# -----------------------------------
# You can now send OpenAI-compatible requests to the proxy. For example:
#
# curl -X POST http://localhost:9000/v1/chat/completions \
# -H "Content-Type: application/json" \
# -d '{
# "model": "your-model",
# "messages": [{"role": "user","content": [{"type": "image_url","image_url": {"url": f"file://{image_path}"}},
# {"type": "text","text": "Describe this image."}]}],
# "max_tokens": 16
# }'
#
# Step 4: Health Check
# --------------------
# To check if the proxy is running and see how many backend instances are
# connected, use:
#
# curl http://localhost:9000/healthcheck
#
# This will return a JSON object with the status and the number of encoder, pd, prefiller
# and decoder instances.
#
# Notes:
# - You can scale the number of encoder, pd, prefiller and decoder servers as needed.
# - The proxy dispatches requests based on a least-loaded strategy,
# using a priority queue to balance the active token workload across instances.
# - For production, ensure your backend servers are robust and secure.
#
# For more details, see the code and comments in this file.
import argparse
import asyncio
import base64
import functools
import heapq
import io
import ipaddress
import math
import os
import sys
import uuid
from contextlib import asynccontextmanager
import httpx
from fastapi import FastAPI, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from PIL import Image
from vllm.logger import init_logger
logger = init_logger(__name__)
# Add uvloop for faster event loop if available
try:
import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
except ImportError:
pass
class ServerState:
def __init__(self, host, port):
self.host = host
self.port = port
self.url = f"http://{host}:{port}/v1"
try:
ip = ipaddress.ip_address(self.host)
if isinstance(ip, ipaddress.IPv6Address):
self.url = f"http://[{host}]:{port}/v1"
except Exception:
pass
self.client = httpx.AsyncClient(
timeout=None,
base_url=self.url,
limits=httpx.Limits(max_connections=100000, max_keepalive_connections=100000),
)
self.active_tokens = 0
self.active_kv_cache = 0
self.active_requests = 0
self.aborted_requests = set()
class ProxyState:
def __init__(self, prefiller_instances, decoder_instances, encoder_instances=None, pd_instances=None):
self.prefillers: list[ServerState] = [ServerState(h, p) for h, p in prefiller_instances]
self.decoders: list[ServerState] = [ServerState(h, p) for h, p in decoder_instances]
self.encoders: list[ServerState] = [ServerState(h, p) for h, p in (encoder_instances or [])]
self.pds: list[ServerState] = [ServerState(h, p) for h, p in pd_instances]
self.req_to_prefiller = {}
self.req_id_lock = asyncio.Lock()
self.prefiller_heap = [(0, i, server) for i, server in enumerate(self.prefillers)]
self.decoder_heap = [(0, i, server) for i, server in enumerate(self.decoders)]
self.encoder_heap = [(0, i, server) for i, server in enumerate(self.encoders)]
self.pd_heap = [(0, i, server) for i, server in enumerate(self.pds)]
heapq.heapify(self.prefiller_heap)
heapq.heapify(self.decoder_heap)
heapq.heapify(self.encoder_heap)
heapq.heapify(self.pd_heap)
self.req_id_future = {}
self.req_data_dict = {}
def _update_pd_priority(self, server_idx: int):
server = self.pds[server_idx]
priority = server.active_tokens + server.active_kv_cache * 0.3
self.pd_heap = [(p, i, s) for p, i, s in self.pd_heap if i != server_idx]
heapq.heappush(self.pd_heap, (priority, server_idx, server)) # type: ignore[misc]
def _update_prefiller_priority(self, server_idx: int):
server = self.prefillers[server_idx]
priority = server.active_tokens + server.active_kv_cache * 0.3
self.prefiller_heap = [(p, i, s) for p, i, s in self.prefiller_heap if i != server_idx]
heapq.heappush(self.prefiller_heap, (priority, server_idx, server)) # type: ignore[misc]
def _update_decoder_priority(self, server_idx: int):
server = self.decoders[server_idx]
priority = server.active_tokens
self.decoder_heap = [(p, i, s) for p, i, s in self.decoder_heap if i != server_idx]
heapq.heappush(self.decoder_heap, (priority, server_idx, server))
def _update_encoder_priority(self, server_idx: int):
server = self.encoders[server_idx]
priority = server.active_tokens
self.encoder_heap = [(p, i, s) for p, i, s in self.encoder_heap if i != server_idx]
heapq.heappush(self.encoder_heap, (priority, server_idx, server))
def abort_pd_request(self, server_idx: int, request_id):
self.pds[server_idx].aborted_requests.add(request_id)
def acquire_aborted_pd_requests(self, server_idx: int):
aborted_requests = self.pds[server_idx].aborted_requests.copy()
self.pds[server_idx].aborted_requests.clear()
return aborted_requests
def abort_prefiller_request(self, server_idx: int, request_id):
self.prefillers[server_idx].aborted_requests.add(request_id)
def acquire_aborted_prefiller_requests(self, server_idx: int):
aborted_requests = self.prefillers[server_idx].aborted_requests.copy()
self.prefillers[server_idx].aborted_requests.clear()
return aborted_requests
async def next_req_id(self):
async with self.req_id_lock:
return str(uuid.uuid4())
def select_pd(self, token_count):
if not self.pd_heap:
raise RuntimeError("No pd servers available")
priority, chosen, server = heapq.heappop(self.pd_heap)
self.pds[chosen].active_tokens += token_count
self.pds[chosen].active_kv_cache += token_count
self._update_pd_priority(chosen)
return chosen
def release_pd(self, idx, token_count):
self.pds[idx].active_tokens -= token_count
self._update_pd_priority(idx)
def select_prefiller(self, token_count):
if not self.prefiller_heap:
raise RuntimeError("No prefiller servers available")
priority, chosen, server = heapq.heappop(self.prefiller_heap)
self.prefillers[chosen].active_tokens += token_count
self.prefillers[chosen].active_kv_cache += token_count
self._update_prefiller_priority(chosen)
return chosen
def release_prefiller(self, idx, token_count):
self.prefillers[idx].active_tokens -= token_count
self._update_prefiller_priority(idx)
def release_prefiller_kv(self, idx, token_count):
if self.prefillers[idx].active_kv_cache > 0:
self.prefillers[idx].active_kv_cache -= token_count
self._update_prefiller_priority(idx)
def select_decoder(self, token_count):
if not self.decoder_heap:
raise RuntimeError("No decoder servers available")
priority, chosen, server = heapq.heappop(self.decoder_heap)
self.decoders[chosen].active_tokens += token_count
self._update_decoder_priority(chosen)
return chosen
def release_decoder(self, idx, token_count):
self.decoders[idx].active_tokens -= token_count
self._update_decoder_priority(idx)
def select_encoder(self, token_count):
if not self.encoder_heap:
raise RuntimeError("No encoder servers available")
priority, chosen, server = heapq.heappop(self.encoder_heap)
self.encoders[chosen].active_tokens += token_count
self._update_encoder_priority(chosen)
return chosen
def release_encoder(self, idx, token_count):
self.encoders[idx].active_tokens -= token_count
self._update_encoder_priority(idx)
def calculate_prefill_scores(self, text_length: int) -> float:
length_score = text_length / 4.0
input_score = length_score * 0.0345 + 120.0745
return input_score
def calculate_decode_scores(self, request_length: int) -> float:
return request_length
proxy_state = None
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--host", type=str, default="localhost")
parser.add_argument("--prefiller-hosts", type=str, nargs="+", default=[])
parser.add_argument("--prefiller-ports", type=int, nargs="+", default=[])
parser.add_argument("--decoder-hosts", type=str, nargs="+", default=[])
parser.add_argument("--decoder-ports", type=int, nargs="+", default=[])
parser.add_argument("--encoder-hosts", type=str, nargs="+", default=[])
parser.add_argument("--encoder-ports", type=int, nargs="+", default=[])
parser.add_argument("--pd-hosts", type=str, nargs="+", default=[])
parser.add_argument("--pd-ports", type=int, nargs="+", default=[])
parser.add_argument("--max-retries", type=int, default=3, help="Maximum number of retries for HTTP requests")
parser.add_argument(
"--retry-delay", type=float, default=0.001, help="Base delay (seconds) for exponential backoff retries"
)
args = parser.parse_args()
if len(args.pd_hosts) != len(args.pd_ports):
raise ValueError("Number of pd hosts must match number of pd ports")
if len(args.prefiller_hosts) != len(args.prefiller_ports):
raise ValueError("Number of prefiller hosts must match number of prefiller ports")
if len(args.decoder_hosts) != len(args.decoder_ports):
raise ValueError("Number of decoder hosts must match number of decoder ports")
if len(args.encoder_hosts) != len(args.encoder_ports):
raise ValueError("Number of encoder hosts must match number of encoder ports")
args.prefiller_instances = list(zip(args.prefiller_hosts, args.prefiller_ports))
args.decoder_instances = list(zip(args.decoder_hosts, args.decoder_ports))
args.encoder_instances = list(zip(args.encoder_hosts, args.encoder_ports))
args.pd_instances = list(zip(args.pd_hosts, args.pd_ports))
return args
@asynccontextmanager
async def lifespan(app: FastAPI):
global proxy_state
proxy_state = ProxyState(
global_args.prefiller_instances,
global_args.decoder_instances,
global_args.encoder_instances,
global_args.pd_instances,
)
print(
f"Initialized {len(proxy_state.encoders)} encode clients, {len(proxy_state.prefillers)} prefill clients \
and \{len(proxy_state.decoders)} decode clients, {len(proxy_state.pds)} pd clients."
)
yield
for e in proxy_state.encoders:
await e.client.aclose()
for p in proxy_state.prefillers:
await p.client.aclose()
for d in proxy_state.decoders:
await d.client.aclose()
for pd in proxy_state.pds:
await pd.client.aclose()
async def listen_for_disconnect(request: Request) -> None:
while True:
message = await request.receive()
if message["type"] == "http.disconnect":
break
def with_cancellation(handler_func):
@functools.wraps(handler_func)
async def wrapper(*args, **kwargs):
request = kwargs["request"]
handler_task = asyncio.create_task(handler_func(*args, **kwargs))
cancellation_task = asyncio.create_task(listen_for_disconnect(request))
done, pending = await asyncio.wait([handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED)
for task in pending:
task.cancel()
if handler_task in done:
return handler_task.result()
return None
return wrapper
app = FastAPI(lifespan=lifespan)
async def send_request_to_encode_service(
client: httpx.AsyncClient,
encoder_id: int,
endpoint: str,
req_data: dict,
request_id: str,
max_retries: int = 3,
base_delay: float = 0.2,
):
encoder_req = req_data.copy()
encoder_req["stream"] = False
encoder_req["max_tokens"] = 1
encoder_req["min_tokens"] = 1
if "stream_options" in encoder_req:
del encoder_req["stream_options"]
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
for attempt in range(1, max_retries + 1):
try:
response = await client.post(endpoint, json=encoder_req, headers=headers)
response.raise_for_status()
return response
except (httpx.RequestError, httpx.HTTPStatusError) as e:
logger.warning("Attempt %s failed for %s: %s", attempt, endpoint, e)
last_exc = e
if attempt < max_retries:
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for %s.", max_retries, endpoint)
raise last_exc
async def stream_service_response_with_retry(
client: httpx.AsyncClient,
endpoint: str,
req_data: dict,
request_id: str,
max_retries: int = 3,
base_delay: float = 0.2,
):
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
for attempt in range(1, max_retries + 1):
try:
async with client.stream("POST", endpoint, json=req_data, headers=headers) as response:
response.raise_for_status()
first_chunk_sent = False
async for chunk in response.aiter_bytes():
first_chunk_sent = True
yield chunk
return
except (httpx.RequestError, httpx.HTTPStatusError) as e:
if attempt < max_retries:
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, e)
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
raise e
except Exception as e:
if "first_chunk_sent" in locals() and first_chunk_sent:
logger.error("Streaming to client interrupted after response started: %s", e)
return
else:
if attempt < max_retries:
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, e)
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
raise e
def fast_get_hw(b64_str):
img_bytes = base64.b64decode(b64_str.split(",")[1])
img = Image.open(io.BytesIO(img_bytes))
return img.width, img.height
SOF_MARKERS = {
0xC0,
0xC1,
0xC2,
0xC3,
0xC5,
0xC6,
0xC7,
0xC9,
0xCA,
0xCB,
0xCD,
0xCE,
0xCF,
}
def parse_jpeg_size(data: bytes):
idx = 0
length = len(data)
if length < 2 or data[0:2] != b"\xff\xd8":
raise ValueError("Not a JPEG")
idx = 2
while idx + 9 < length:
if data[idx] != 0xFF:
idx += 1
continue
marker = data[idx + 1]
# 跳过填充字节
if marker == 0xFF:
idx += 1
continue
if marker in SOF_MARKERS:
h = (data[idx + 5] << 8) | data[idx + 6]
w = (data[idx + 7] << 8) | data[idx + 8]
return w, h
if marker in (0xD9, 0xDA):
break
if idx + 3 >= length:
break
seg_len = (data[idx + 2] << 8) | data[idx + 3]
if seg_len < 2:
break
idx += 2 + seg_len
raise ValueError("JPEG SOF marker not found")
def parse_png_size(data: bytes):
w = int.from_bytes(data[16:20], "big")
h = int.from_bytes(data[20:24], "big")
return w, h
def get_hw_from_local(path: str):
if path.startswith("file://"):
path = path[7:]
with open(path, "rb") as f:
data = f.read(65536)
if data.startswith(b"\x89PNG"):
return parse_png_size(data)
return parse_jpeg_size(data)
def calculate_messages_size(ori_req_data, ori_req_body):
messages = ori_req_data.get("messages")
stats = {
"text_char_count": 0,
"mul_token": 0,
}
for msg in messages:
if not isinstance(msg.get("content"), list):
continue
for content_item in msg["content"]:
content_type = content_item.get("type")
if not content_type:
continue
if content_type == "text":
text = content_item.get("text", "")
stats["text_char_count"] += len(text)
elif content_type == "image_url":
img_url = content_item.get("image_url", {}).get("url", "")
if img_url.startswith("data:image"):
h, w = fast_get_hw(img_url)
else:
h, w = get_hw_from_local(img_url)
stats["mul_token"] += math.ceil(h / 32) * math.ceil(w / 32)
elif content_type == "video_url":
stats["mul_token"] += len(ori_req_body) * 32
return stats
def get_api_request_id(api, req_id):
if api == "/completions":
return "cmpl-" + req_id + "-0"
elif api == "/chat/completions":
return "chatcmpl-" + req_id
def get_origin_request_id(api, req_id):
if api == "/completions":
return req_id.replace("cmpl-", "")[:-2]
elif api == "/chat/completions":
return req_id.replace("chatcmpl-", "")
async def non_stream_retry_wrap(forward_func, max_retries: int = 3, delay: float = 0.001):
last_exc = None
for attempt in range(max_retries):
try:
result = await forward_func()
return result
except Exception as e:
if isinstance(e, HTTPException) and e.status_code < 500:
raise
last_exc = e
logger.warning(
"attempt %s / %s failed retrying... ",
attempt + 1,
max_retries,
)
await asyncio.sleep(delay * (attempt + 1))
raise RuntimeError(f"all {max_retries} retries failed.") from last_exc
async def _handle_completions(api: str, request: Request):
try:
req_data = await request.json()
req_body = await request.body()
request_id = await proxy_state.next_req_id()
request_id_api = get_api_request_id(api, request_id)
mul_flag = False
stats_info = calculate_messages_size(req_data, req_body)
text_length = stats_info["text_char_count"]
encoder_score = stats_info["mul_token"]
if stats_info["mul_token"] != 0:
mul_flag = True
if mul_flag and proxy_state.encoders:
encoder_idx = proxy_state.select_encoder(encoder_score)
encoder = proxy_state.encoders[encoder_idx]
logger.debug("Sending to encoder: %s", encoder.url)
_ = await send_request_to_encode_service(
encoder.client,
encoder_idx,
api,
req_data,
request_id,
max_retries=global_args.max_retries,
base_delay=global_args.retry_delay,
)
proxy_state.release_encoder(encoder_idx, encoder_score)
token_score = encoder_score + text_length
if proxy_state.pds:
pd_idx = proxy_state.select_pd(token_score)
pd = proxy_state.pds[pd_idx]
async def generate_stream():
try:
async for chunk in stream_service_response_with_retry(
pd.client,
api,
req_data,
request_id=request_id,
max_retries=global_args.max_retries,
base_delay=global_args.retry_delay,
):
yield chunk
except Exception as e:
logger.error("Error during streaming from pd %s: %s", pd.url, e)
proxy_state.abort_pd_request(pd_idx, request_id)
finally:
proxy_state.release_pd(pd_idx, token_score)
return StreamingResponse(generate_stream(), media_type="application/json")
else:
proxy_state.req_data_dict[request_id_api] = (req_data, token_score, api)
req_data["kv_transfer_params"] = {
"do_remote_decode": False,
"do_remote_prefill": True,
"metaserver": f"http://{global_args.host}:{global_args.port}/v1/metaserver",
}
# Select decoder
decoder_score = proxy_state.calculate_decode_scores(token_score)
logger.debug("Decoder score: %f", decoder_score)
# Use the prefiller's kv_transfer_params to select decoder
decoder_idx = proxy_state.select_decoder(decoder_score)
print("d", decoder_idx, decoder_score)
decoder = proxy_state.decoders[decoder_idx]
# logger.debug("Using %s %s", prefiller.url, decoder.url)
# Stream response from decoder
released_kv = False
async def generate_stream():
nonlocal released_kv
try:
async for chunk in stream_service_response_with_retry(
decoder.client,
api,
req_data,
request_id=request_id,
max_retries=global_args.max_retries,
base_delay=global_args.retry_delay,
):
yield chunk
except Exception as e:
logger.error(
"Error during streaming from decoder %s: %s the aborted request %s "
"will be routing to the target prefiller when new request is ready to dispatch to it",
decoder.url,
e,
request_id,
)
# After streaming done, release tokens
proxy_state.release_decoder(decoder_idx, decoder_score)
return StreamingResponse(generate_stream(), media_type="application/json")
except Exception as e:
import traceback
exc_info = sys.exc_info()
print(f"Error occurred in disagg prefill proxy server - {api} endpoint")
print(e)
print("".join(traceback.format_exception(*exc_info)))
raise
@app.post("/v1/completions")
@with_cancellation
async def handle_completions(request: Request):
return await _handle_completions("/completions", request)
@app.post("/v1/chat/completions")
@with_cancellation
async def handle_chat_completions(request: Request):
return await _handle_completions("/chat/completions", request)
@app.get("/healthcheck")
async def healthcheck():
return {
"status": "ok",
"encode_instances": len(proxy_state.encoders),
"prefill_instances": len(proxy_state.prefillers),
"decode_instances": len(proxy_state.decoders),
"pd_instances": len(proxy_state.pds),
}
async def send_request_to_service(
client: httpx.AsyncClient,
prefiller_id: int,
endpoint: str,
req_data: dict,
request_id: str,
max_retries: int = 3,
base_delay: float = 0.2,
):
req_data = req_data.copy()
req_data["stream"] = False
req_data["max_tokens"] = 1
req_data["min_tokens"] = 1
if "max_completion_tokens" in req_data:
req_data["max_completion_tokens"] = 1
if "stream_options" in req_data:
del req_data["stream_options"]
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
last_exc = None
for attempt in range(1, max_retries + 1):
try:
response = await client.post(endpoint, json=req_data, headers=headers)
response.raise_for_status()
if request_id in proxy_state.req_id_future:
result_future = proxy_state.req_id_future[request_id]
result_future.set_result(response.json()["kv_transfer_params"])
return
except (httpx.RequestError, httpx.HTTPStatusError) as e:
logger.warning("Attempt %s failed for %s: %s", attempt, endpoint, e)
last_exc = e
if attempt < max_retries:
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for %s.", max_retries, endpoint)
raise last_exc
@app.post("/v1/metaserver")
async def metaserver(request: Request):
try:
kv_transfer_params = await request.json()
request_id = kv_transfer_params["request_id"]
assert request_id in proxy_state.req_data_dict
req_data, token_score, api = proxy_state.req_data_dict[request_id]
request_id = get_origin_request_id(api, request_id)
req_data["kv_transfer_params"] = kv_transfer_params
logger.debug("Prefiller score: %s", token_score)
# Select prefiller
prefiller_idx = proxy_state.select_prefiller(token_score)
prefiller = proxy_state.prefillers[prefiller_idx]
logger.debug("Using prefill prefiller.url=%r req_data=%r", prefiller.url, req_data)
# Send request to prefiller
_ = await send_request_to_service(
prefiller.client,
prefiller_idx,
api,
req_data,
request_id,
max_retries=global_args.max_retries,
base_delay=global_args.retry_delay,
)
proxy_state.release_prefiller(prefiller_idx, token_score)
proxy_state.release_prefiller_kv(prefiller_idx, token_score)
except Exception as e:
logger.error("Post metaserver failed with: %s", e)
proxy_state.release_prefiller(prefiller_idx, token_score)
proxy_state.release_prefiller_kv(prefiller_idx, token_score)
###################################### profile ######################################
async def _forward_profile(
service_name: str, idx: int, client, host: str, port: int, endpoint: str, req_data: dict, headers: dict
):
"""Forward profiling request to one service and return raw response or error."""
url = f"http://{host}:{port}{endpoint}"
try:
resp = await client.post(url, json=req_data, headers=headers, timeout=10.0)
resp.raise_for_status()
# 直接返回 httpx.Response，保持原始格式
return service_name, idx, {"status_code": resp.status_code, "body": resp.text}
except Exception as e:
return service_name, idx, {"error": str(e)}
@app.post("/start_profile")
async def start_profile(request: Request):
"""
Forward the /start_profile request to all encoder, prefiller, and decoder services (concurrently).
"""
try:
try:
req_data = await request.json()
except Exception as e:
print(f"Error in stop_profile while waiting request data: {e}")
req_data = {}
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
tasks = []
# encoder
for idx, encoder in enumerate(proxy_state.encoders):
tasks.append(
_forward_profile(
"encoder", idx, encoder.client, encoder.host, encoder.port, "/start_profile", req_data, headers
)
)
# prefiller
for idx, prefill in enumerate(proxy_state.prefillers):
tasks.append(
_forward_profile(
"prefill", idx, prefill.client, prefill.host, prefill.port, "/start_profile", req_data, headers
)
)
# decoder
for idx, decoder in enumerate(proxy_state.decoders):
tasks.append(
_forward_profile(
"decoder", idx, decoder.client, decoder.host, decoder.port, "/start_profile", req_data, headers
)
)
# pds
for idx, pd in enumerate(proxy_state.pds):
tasks.append(_forward_profile("pds", idx, pd.client, pd.host, pd.port, "/start_profile", req_data, headers))
results_list = await asyncio.gather(*tasks)
results = {f"{name}_{idx}": res for name, idx, res in results_list}
return JSONResponse(content={"status": "done", "results": results}, status_code=200)
except Exception as e:
print(f"Error in start_profile: {e}")
return JSONResponse(content={"error": str(e)}, status_code=500)
@app.post("/stop_profile")
async def stop_profile(request: Request):
"""
Forward the /stop_profile request to all encoder, prefiller, and decoder services (concurrently).
"""
try:
try:
req_data = await request.json()
except Exception as e:
print(f"Error in stop_profile while waiting request data: {e}")
req_data = {}
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}"}
tasks = []
# encoder
for idx, encoder in enumerate(proxy_state.encoders):
tasks.append(
_forward_profile(
"encoder", idx, encoder.client, encoder.host, encoder.port, "/stop_profile", req_data, headers
)
)
# prefiller
for idx, prefill in enumerate(proxy_state.prefillers):
tasks.append(
_forward_profile(
"prefill", idx, prefill.client, prefill.host, prefill.port, "/stop_profile", req_data, headers
)
)
# decoder
for idx, decoder in enumerate(proxy_state.decoders):
tasks.append(
_forward_profile(
"decoder", idx, decoder.client, decoder.host, decoder.port, "/stop_profile", req_data, headers
)
)
# pds
for idx, pd in enumerate(proxy_state.pds):
tasks.append(_forward_profile("pds", idx, pd.client, pd.host, pd.port, "/stop_profile", req_data, headers))
results_list = await asyncio.gather(*tasks)
results = {f"{name}_{idx}": res for name, idx, res in results_list}
return JSONResponse(content={"status": "done", "results": results}, status_code=200)
except Exception as e:
print(f"Error in stop_profile: {e}")
return JSONResponse(content={"error": str(e)}, status_code=500)
if __name__ == "__main__":
global global_args
global_args = parse_args()
import uvicorn
uvicorn.run(app, host=global_args.host, port=global_args.port)

107
examples/eplb/eplb_deepseek.py
View File

@@ -4,13 +4,11 @@ Expert parallelism load balancer (EPLB) for vLLM.
The rearrangement algorithm is adapted from
[DeepSeek EPLB](https://github.com/deepseek-ai/eplb).
"""
from typing import Tuple
import torch
def balanced_packing(weight: torch.Tensor,
num_packs: int) -> Tuple[torch.Tensor, torch.Tensor]:
def balanced_packing(weight: torch.Tensor, num_packs: int) -> tuple[torch.Tensor, torch.Tensor]:
"""
Pack n weighted objects to m packs, such that each bin contains exactly n/m objects and the weights of all packs
are as balanced as possible.
@@ -18,8 +16,8 @@ def balanced_packing(weight: torch.Tensor,
Parameters:
weight: [X, n], the weight of each item
num_packs: number of packs
Returns:
Returns:
pack_index: [X, n], the pack index of each item
rank_in_pack: [X, n], the rank of the item in the pack
"""
@@ -28,26 +26,18 @@ def balanced_packing(weight: torch.Tensor,
groups_per_pack = num_groups // num_packs
if groups_per_pack == 1:
pack_index = torch.arange(weight.size(-1),
dtype=torch.int64,
device=weight.device).expand(weight.shape)
pack_index = torch.arange(weight.size(-1), dtype=torch.int64, device=weight.device).expand(weight.shape)
rank_in_pack = torch.zeros_like(weight, dtype=torch.int64)
return pack_index, rank_in_pack
indices = weight.float().sort(-1, descending=True).indices.cpu()
pack_index = torch.full_like(weight,
fill_value=-1,
dtype=torch.int64,
device='cpu')
pack_index = torch.full_like(weight, fill_value=-1, dtype=torch.int64, device="cpu")
rank_in_pack = torch.full_like(pack_index, fill_value=-1)
for i in range(num_layers):
pack_weights = [0] * num_packs
pack_items = [0] * num_packs
for group in indices[i]:
pack = min(
(i
for i in range(num_packs) if pack_items[i] < groups_per_pack),
key=pack_weights.__getitem__)
pack = min((i for i in range(num_packs) if pack_items[i] < groups_per_pack), key=pack_weights.__getitem__)
assert pack_items[pack] < groups_per_pack
pack_index[i, group] = pack
rank_in_pack[i, group] = pack_items[pack]
@@ -56,16 +46,14 @@ def balanced_packing(weight: torch.Tensor,
return pack_index, rank_in_pack
def replicate_experts(
weight: torch.Tensor,
num_phy: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
def replicate_experts(weight: torch.Tensor, num_phy: int) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Replicate `num_log` experts to `num_phy` replicas, such that the maximum load of all replicas is minimized.
Parameters:
weight: [X, num_log]
num_phy: total number of experts after replication
Returns:
phy2log: [X, num_phy], logical expert id of each physical expert
rank: [X, num_phy], the replica rank
@@ -75,8 +63,7 @@ def replicate_experts(
num_redundant = num_phy - num_log
assert num_redundant >= 0
device = weight.device
phy2log = torch.arange(num_phy, dtype=torch.int64,
device=device).repeat(n, 1)
phy2log = torch.arange(num_phy, dtype=torch.int64, device=device).repeat(n, 1)
rank = torch.zeros(n, num_phy, dtype=torch.int64, device=device)
logcnt = torch.ones(n, num_log, dtype=torch.int64, device=device)
arangen = torch.arange(n, dtype=torch.int64, device=device)
@@ -88,9 +75,9 @@ def replicate_experts(
return phy2log, rank, logcnt
def rebalance_experts_hierarchical(weight: torch.Tensor,
num_physical_experts: int, num_groups: int,
num_nodes: int, num_gpus: int):
def rebalance_experts_hierarchical(
weight: torch.Tensor, num_physical_experts: int, num_groups: int, num_nodes: int, num_gpus: int
):
"""
Parameters:
weight: [num_moe_layers, num_logical_experts]
@@ -99,7 +86,7 @@ def rebalance_experts_hierarchical(weight: torch.Tensor,
num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
num_gpus: number of GPUs, must be a multiple of `num_nodes`
Returns:
Returns:
physical_to_logical_map: [num_moe_layers, num_physical_experts]
logical_to_physical_map: [num_moe_layers, num_logical_experts, X]
logical_count: [num_moe_layers, num_logical_experts]
@@ -115,45 +102,37 @@ def rebalance_experts_hierarchical(weight: torch.Tensor,
def inverse(perm: torch.Tensor) -> torch.Tensor:
inv = torch.empty_like(perm)
inv.scatter_(
1, perm,
torch.arange(perm.size(1), dtype=torch.int64,
device=perm.device).expand(perm.shape))
inv.scatter_(1, perm, torch.arange(perm.size(1), dtype=torch.int64, device=perm.device).expand(perm.shape))
return inv
# Step 1: pack groups to nodes
tokens_per_group = weight.unflatten(-1, (num_groups, group_size)).sum(-1)
group_pack_index, group_rank_in_pack = balanced_packing(
tokens_per_group, num_nodes)
log2mlog = (((group_pack_index * groups_per_node + group_rank_in_pack) *
group_size).unsqueeze(-1) +
torch.arange(group_size,
dtype=torch.int64,
device=group_pack_index.device)).flatten(-2)
group_pack_index, group_rank_in_pack = balanced_packing(tokens_per_group, num_nodes)
log2mlog = (
((group_pack_index * groups_per_node + group_rank_in_pack) * group_size).unsqueeze(-1)
+ torch.arange(group_size, dtype=torch.int64, device=group_pack_index.device)
).flatten(-2)
mlog2log = inverse(log2mlog)
# Step 2: construct redundant experts within nodes
# [num_layers * num_nodes, num_logical_experts // num_nodes]
tokens_per_mlog = weight.gather(-1, mlog2log).view(
-1, num_logical_experts // num_nodes)
phy2mlog, phyrank, mlogcnt = replicate_experts(
tokens_per_mlog, num_physical_experts // num_nodes)
tokens_per_mlog = weight.gather(-1, mlog2log).view(-1, num_logical_experts // num_nodes)
phy2mlog, phyrank, mlogcnt = replicate_experts(tokens_per_mlog, num_physical_experts // num_nodes)
# Step 3: pack physical_experts to GPUs
# [num_layers * num_nodes, num_physical_experts // num_nodes]
tokens_per_phy = (tokens_per_mlog / mlogcnt).gather(-1, phy2mlog)
pack_index, rank_in_pack = balanced_packing(tokens_per_phy,
num_gpus // num_nodes)
pack_index, rank_in_pack = balanced_packing(tokens_per_phy, num_gpus // num_nodes)
phy2pphy = pack_index * phy_experts_per_gpu + rank_in_pack
pphy2phy = inverse(phy2pphy)
pphy2mlog = phy2mlog.gather(
-1, pphy2phy) # [num_layers * num_nodes, num_log_per_nodes]
pphy2mlog = (pphy2mlog.view(num_layers, num_nodes, -1) + torch.arange(
0,
num_logical_experts,
num_logical_experts // num_nodes,
device=group_pack_index.device).view(1, -1, 1)).flatten(-2)
pphy2mlog = phy2mlog.gather(-1, pphy2phy) # [num_layers * num_nodes, num_log_per_nodes]
pphy2mlog = (
pphy2mlog.view(num_layers, num_nodes, -1)
+ torch.arange(0, num_logical_experts, num_logical_experts // num_nodes, device=group_pack_index.device).view(
1, -1, 1
)
).flatten(-2)
pphy2log = mlog2log.gather(-1, pphy2mlog)
pphyrank = phyrank.gather(-1, pphy2phy).view(num_layers, -1)
logcnt = mlogcnt.view(num_layers, -1).gather(-1, log2mlog)
@@ -161,9 +140,8 @@ def rebalance_experts_hierarchical(weight: torch.Tensor,
def rebalance_experts(
weight: torch.Tensor, num_replicas: int, num_groups: int,
num_nodes: int,
num_gpus: int) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
weight: torch.Tensor, num_replicas: int, num_groups: int, num_nodes: int, num_gpus: int
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
"""
Entry point for expert-parallelism load balancer.
@@ -174,7 +152,7 @@ def rebalance_experts(
num_nodes: number of server nodes, where the intra-node network (e.g, NVLink) is faster
num_gpus: number of GPUs, must be a multiple of `num_nodes`
Returns:
Returns:
physical_to_logical_map: [layers, num_replicas], the expert index of each replica
logical_to_physical_map: [layers, num_logical_experts, X], the replica indices for each expert
expert_count: [layers, num_logical_experts], number of physical replicas for each logical expert
@@ -183,23 +161,20 @@ def rebalance_experts(
weight = weight.float().cpu()
if num_groups % num_nodes == 0:
# use hierarchical load-balance policy
phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
weight, num_replicas, num_groups, num_nodes, num_gpus)
phy2log, phyrank, logcnt = rebalance_experts_hierarchical(weight, num_replicas, num_groups, num_nodes, num_gpus)
else:
# use global load-balance policy
phy2log, phyrank, logcnt = rebalance_experts_hierarchical(
weight, num_replicas, 1, 1, num_gpus)
phy2log, phyrank, logcnt = rebalance_experts_hierarchical(weight, num_replicas, 1, 1, num_gpus)
maxlogcnt = logcnt.max().item()
log2phy: torch.Tensor = torch.full(
(num_layers, num_logical_experts, maxlogcnt),
-1,
dtype=torch.int64,
device=logcnt.device)
(num_layers, num_logical_experts, maxlogcnt), -1, dtype=torch.int64, device=logcnt.device
)
log2phy.view(num_layers, -1).scatter_(
-1, phy2log * maxlogcnt + phyrank,
torch.arange(num_replicas, dtype=torch.int64,
device=log2phy.device).expand(num_layers, -1))
-1,
phy2log * maxlogcnt + phyrank,
torch.arange(num_replicas, dtype=torch.int64, device=log2phy.device).expand(num_layers, -1),
)
return phy2log, log2phy, logcnt
__all__ = ['rebalance_experts']
__all__ = ["rebalance_experts"]

79
examples/eplb/eplb_strategy.py
View File

@@ -1,7 +1,5 @@
# coding=utf-8
# Copyright (c) Huawei Technologies Co., Ltd. 2025-2025. All rights reserved.
import json
import logging
import os
import matplotlib.pyplot as plt # type: ignore
@@ -11,8 +9,6 @@ import torch
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
logger = logging.getLogger("msit_logger")
def save_matrix_to_json(output_path, file_name, deployment):
num_layers = deployment.shape[0]
@@ -24,10 +20,7 @@ def save_matrix_to_json(output_path, file_name, deployment):
layer = {"layer_id": i, "device_count": num_cards}
device_list = []
for j in range(num_cards):
device = {
"device_id": j,
"device_expert": deployment[i, j].tolist()
}
device = {"device_id": j, "device_expert": deployment[i, j].tolist()}
device_list.append(device)
layer["device_list"] = device_list
layer_list.append(layer)
@@ -37,7 +30,7 @@ def save_matrix_to_json(output_path, file_name, deployment):
# Save as JSON file
try:
with open(file_name, 'w') as f:
with open(file_name, "w") as f:
json.dump(data, f, indent=4)
except Exception as e:
print(f"write {file_name} failed: {e}")
@@ -66,21 +59,17 @@ def calculate_average(lst):
return total / count
def layer_imblance_polt(y_list, label_names, device_num, output_path,
file_name):
plt.rcParams['font.sans-serif'] = ['Arial']
plt.rcParams['axes.unicode_minus'] = False
def layer_imbalance_plot(y_list, label_names, device_num, output_path, file_name):
plt.rcParams["font.sans-serif"] = ["Arial"]
plt.rcParams["axes.unicode_minus"] = False
x = [i for i in range(58)]
for index, y in enumerate(y_list):
plt.plot(x,
y,
label=rf'{label_names[index]}，avg={calculate_average(y)}')
plt.plot(x, y, label=rf"{label_names[index]}，avg={calculate_average(y)}")
plt.legend()
plt.title(rf'Load Distribution (num_gpus={device_num})')
plt.xlabel('layer')
plt.ylabel('Device Load')
plt.title(rf"Load Distribution (num_gpus={device_num})")
plt.xlabel("layer")
plt.ylabel("Device Load")
# Show grid lines
plt.grid(True)
@@ -91,27 +80,23 @@ def layer_imblance_polt(y_list, label_names, device_num, output_path,
plt.close()
def deepseek_deploy(workload, num_redundancy_expert, num_groups, num_nodes,
num_gpus, num_original_expert):
def deepseek_deploy(workload, num_redundancy_expert, num_groups, num_nodes, num_gpus, num_original_expert):
from eplb_deepseek import rebalance_experts
num_replicas = num_original_expert + num_redundancy_expert
hy2log, log2phy, logcnt = rebalance_experts(workload, num_replicas,
num_groups, num_nodes,
num_gpus)
hy2log, log2phy, logcnt = rebalance_experts(workload, num_replicas, num_groups, num_nodes, num_gpus)
# Convert to global_deployment
workload = workload.cpu().numpy()
global_deployment = []
layer_num = log2phy.shape[0]
num_physical_experts_local = (num_original_expert +
num_redundancy_expert) // num_gpus
num_physical_experts_local = (num_original_expert + num_redundancy_expert) // num_gpus
for layer_idx in range(layer_num):
layer_deployment = []
for gpu_idx in range(num_gpus):
local_deployment = hy2log[layer_idx][gpu_idx *
num_physical_experts_local:
(gpu_idx + 1) *
num_physical_experts_local]
local_deployment = hy2log[layer_idx][
gpu_idx * num_physical_experts_local : (gpu_idx + 1) * num_physical_experts_local
]
local_deployment = local_deployment.flatten()
layer_deployment.append(local_deployment.tolist())
global_deployment.append(layer_deployment)
@@ -125,18 +110,15 @@ def deepseek_deploy(workload, num_redundancy_expert, num_groups, num_nodes,
new_value = workload[layer_idx].reshape(num_gpus, -1)
row_sum = np.sum(new_value, axis=1)
original_weights.append(row_sum.max())
average_weights.append((np.sum(workload[layer_idx]) / num_gpus))
average_weights.append(np.sum(workload[layer_idx]) / num_gpus)
opt_workload = np.zeros((num_original_expert + num_redundancy_expert),
dtype=np.float64)
opt_workload = np.zeros((num_original_expert + num_redundancy_expert), dtype=np.float64)
for expert_idx in range(num_original_expert):
physical_expert_idxs = log2phy[layer_idx][expert_idx]
physical_expert_idxs = physical_expert_idxs.flatten()
physical_expert_idxs = physical_expert_idxs[
physical_expert_idxs != -1]
physical_expert_idxs = physical_expert_idxs[physical_expert_idxs != -1]
for physical_expert_idx in physical_expert_idxs:
opt_workload[physical_expert_idx] += workload[layer_idx][
expert_idx] / len(physical_expert_idxs)
opt_workload[physical_expert_idx] += workload[layer_idx][expert_idx] / len(physical_expert_idxs)
opt_workload = opt_workload.reshape(num_gpus, -1)
row_sum = np.sum(opt_workload, axis=1)
max_weights.append(row_sum.max())
@@ -145,8 +127,9 @@ def deepseek_deploy(workload, num_redundancy_expert, num_groups, num_nodes,
return global_deployment, y_list
if __name__ == '__main__':
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser()
parser.add_argument("--exp_name", type=str, default="gsm8k_temp0.0")
parser.add_argument("--num_original_expert", type=int, default=256)
@@ -168,19 +151,13 @@ if __name__ == '__main__':
num_nodes = args.num_nodes
# NOTE: assume input workload format: [layer_num, num_experts]
workload = torch.load(input_path, map_location=torch.device('cpu'))
global_deployment, y_list = deepseek_deploy(workload,
num_redundancy_expert,
num_groups, num_nodes,
num_devices,
num_original_expert)
workload = torch.load(input_path, map_location=torch.device("cpu"))
global_deployment, y_list = deepseek_deploy(
workload, num_redundancy_expert, num_groups, num_nodes, num_devices, num_original_expert
)
file_name = f"{exp_name}_{num_devices}_{num_redundancy_expert}"
save_matrix_to_json(output_path, file_name, np.array(global_deployment))
label_names = [
'default deployment max load', 'balanced load max load',
'balanced load avg load'
]
label_names = ["default deployment max load", "balanced load max load", "balanced load avg load"]
new_file_name = f"{exp_name}_{num_devices}_{num_redundancy_expert}.png"
layer_imblance_polt(y_list, label_names, num_devices, output_path,
new_file_name)
layer_imbalance_plot(y_list, label_names, num_devices, output_path, new_file_name)

27
examples/external_online_dp/README.md
View File

@@ -1,12 +1,14 @@
Here is an example guiding how to use `launch_online_dp.py` to launch external dp server in vllm. User can easily launch external dp server following the steps below:
Here is an example guiding how to use `launch_online_dp.py` to launch external dp vLLM servers. User can easily launch external dp servers following the steps below:
### Modify parameters in `run_dp_template.sh`
`run_dp_template.sh` is an template script used to launch each dp vllm instance separately. It will be called by `launch_online_dp.py` in multi threads and most of its configurations are set by `launch_online_dp.py`. Parameters you need to set manually include:
`run_dp_template.sh` is a template script used to launch each data parallel (dp) vLLM instance separately. It will be called by `launch_online_dp.py` in multiple threads and most of its configurations are set by `launch_online_dp.py`. Parameters you need to set manually include:
1. The IP and socket_ifname of your machine. If running on multi-nodes, please make sure the scripts on each node has been set with correct IP and socket_ifname of that node.
2. vLLM serving related parameters including model_path and other configurations. Note that port, dp-related parammeters and tp_size is set by `launch_online_dp.py`, all the other vLLM parameters in this file only serve as an example and you are free to modify them according to your purpose.
2. vLLM serving related parameters including model_path and other configurations. Note that port, dp-related parameters and tp_size is set by `launch_online_dp.py`, all the other vLLM parameters in this file only serve as an example and you are free to modify them according to your purpose.
### Run `launch_online_dp.py` with CL arguments
All the arguments that can be set by users are:
1. `--dp-size`: global data parallel size, must be set
@@ -18,6 +20,7 @@ All the arguments that can be set by users are:
7. `--vllm-start-port`: Starting port of vLLM serving instances, default 9000
An example of running external DP in one single node:
```(python)
cd examples/external_online_dp
# running DP4 TP4 in a node with 16 NPUs
@@ -25,6 +28,7 @@ python launch_online_dp.py --dp-size 4 --tp-size 4 --dp-size-local 4 --dp-rank-s
```
An example of running external DP in two nodes:
```(python)
cd examples/external_online_dp
# running DP4 TP4 in two nodes with 8 NPUs each
@@ -36,3 +40,20 @@ python launch_online_dp.py --dp-size 4 --tp-size 4 --dp-size-local 2 --dp-rank-s
python launch_online_dp.py --dp-size 4 --tp-size 4 --dp-size-local 2 --dp-rank-start 2 --dp-address x.x.x.x --dp-rpc-port 12342
```
### (Optional) Run `dp_load_balance_proxy_server.py` to load balance requests between external dp servers
External dp server means that you need to handle load balance between multiple dp instances out of vLLM by implementing your custom proxy server. Here we provide an example of request-length-aware dp load-balance proxy server for you. The arguments of `dp_load_balance_proxy_server.py` include:
1. `--port`: port of proxy server, default 8000
2. `--host`: host address of proxy server, default localhost
3. `--dp-hosts`: host addresses of external dp servers
4. `--dp-ports`: ports of external dp servers, the number of dp ports should be the same as dp hosts.
5. `--max-retries`: Max number of retries for HTTP requests, default 3
For example, if you have two external dp servers running in x.x.x.a:10001 and x.x.x.b:10002, then you can start the proxy server by:
```(python)
python dp_load_balance_proxy_server.py --host x.x.x.c --port 8000 --dp-hosts x.x.x.a x.x.x.b --dp-ports 10001 10002
```
which will then serve as the entrypoint for inference requests at x.x.x.c:8000, and load balance coming requests between these two external dp servers according to request length.

371
examples/external_online_dp/dp_load_balance_proxy_server.py Normal file
View File

@@ -0,0 +1,371 @@
# Adapted from https://github.com/vllm-project/vllm/tests/v1/kv_connector/nixl_integration/toy_proxy_server.py
# SPDX-License-Identifier: Apache-2.0
#
# Tutorial: Using the Load Balance Proxy Server For External DP
#
# This proxy server is designed to distribute requests between multiple
# vLLM servers running in data parallel for large language model inference.
# It is useful for scaling out inference workloads and balancing load across
# multiple vLLM instances.
#
# Features:
# - Load balances requests to multiple vLLM servers.
# - Supports OpenAI-compatible /v1/completions and /v1/chat/completions endpoints.
# - Streams responses from backend servers to clients.
#
# Prerequisites:
# - Python 3.10+
# - Install dependencies:
# pip install fastapi<0.124.0 httpx uvicorn
#
# Step 1: Start Your Backend Servers
# ----------------------------------
# You need to have at least two vLLM servers running in data parallel.
# These can be mock servers or actual vLLM servers.
# Note that this proxy also works with only one vLLM server running, but
# will fall back to direct request forwarding which is meaningless.
#
# For testing, you can use the provided mock server:
#
# vllm serve --host 0.0.0.0 --port 8100 --data-parallel-rank 0 ... # vLLM DP0
# vllm serve --host 0.0.0.0 --port 8101 --data-parallel-rank 1 ... # vLLM DP1
#
# Step 2: Start the Proxy Server
# ------------------------------
# Run the proxy server, specifying the host/port for each vLLM DP Instance:
#
# python dp_load_balance_proxy_server.py \
# --host 0.0.0.0 --port 9000 \
# --dp-hosts 127.0.0.1 127.0.0.1 \
# --dp-ports 8100 8101 \
#
# This will start the proxy on port 9000, load balancing between two vLLM DP servers.
#
# Step 3: Send a Request to the Proxy
# -----------------------------------
# You can now send OpenAI-compatible requests to the proxy. For example:
#
# curl -X POST http://localhost:9000/v1/completions \
# -H "Content-Type: application/json" \
# -d '{
# "model": "your-model",
# "prompt": "The quick brown fox jumps over the lazy dog",
# "max_tokens": 16
# }'
#
# Or for chat completions:
#
# curl -X POST http://localhost:9000/v1/chat/completions \
# -H "Content-Type: application/json" \
# -d '{
# "model": "your-model",
# "messages": [{"role": "user", "content": "Hello!"}],
# "max_tokens": 16
# }'
#
# Step 4: Health Check
# --------------------
# To check if the proxy is running and see how many backend instances are
# connected, use:
#
# curl http://localhost:9000/healthcheck
#
# This will return a JSON object with the status and the number of vLLM DP servers.
#
# Notes:
# - You can scale the number of vLLM data parallel size as needed.
# - The proxy will consider the length of requests to balance load.
# - For production, ensure your backend servers are robust and secure.
#
# For more details, see the code and comments in this file.
import argparse
import asyncio
import functools
import heapq
import os
import sys
import uuid
from contextlib import asynccontextmanager
from dataclasses import dataclass
from typing import Any
import httpx
from fastapi import FastAPI, Request
from fastapi.responses import StreamingResponse
from vllm.logger import init_logger
logger = init_logger(__name__)
# Add uvloop for faster event loop if available
try:
import uvloop
asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
except ImportError:
pass
class ServerState:
def __init__(self, host, port):
self.host = host
self.port = port
self.url = f"http://{host}:{port}/v1"
self.client = httpx.AsyncClient(
timeout=None,
base_url=self.url,
limits=httpx.Limits(max_connections=100000, max_keepalive_connections=100000),
)
self.active_tokens = 0
self.aborted_requests = set() # Track aborted requests
class ProxyState:
def __init__(self, server_instances):
self.dp_servers: list[ServerState] = [ServerState(h, p) for h, p in server_instances]
self.req_id_lock = asyncio.Lock()
# Removed selection locks - no longer needed for synchronous methods
# Initialize priority queues for efficient server selection
# Each entry is (priority_score, server_index, server_reference)
# Lower priority score = higher priority (less loaded)
self.lb_heap = [(0, i, server) for i, server in enumerate(self.dp_servers)]
heapq.heapify(self.lb_heap)
def _update_server_priority(self, server_idx: int):
"""Update the priority of a decoder server in the heap."""
server = self.dp_servers[server_idx]
priority = server.active_tokens
# Remove old entry and add new one
self.lb_heap = [(p, i, s) for p, i, s in self.lb_heap if i != server_idx]
heapq.heappush(self.lb_heap, (priority, server_idx, server)) # type: ignore
async def next_req_id(self):
async with self.req_id_lock:
return str(uuid.uuid4())
def select_server(self, token_count): # Changed to synchronous
# No lock needed - entire function is atomic
if not self.lb_heap:
raise RuntimeError("No decoder servers available")
priority, chosen, server = heapq.heappop(self.lb_heap)
# Update the chosen server atomically
self.dp_servers[chosen].active_tokens += token_count
# Update priority and re-add to heap
self._update_server_priority(chosen)
return chosen
def release_server(self, idx: int, token_count): # Changed to synchronous
# No lock needed - atomic operation
self.dp_servers[idx].active_tokens -= token_count
# Update priority queue after releasing
self._update_server_priority(idx)
def calculate_request_score(self, request_length: int, max_tokens: int = 16, ignore_eos: bool = False) -> float:
if ignore_eos:
return request_length + max_tokens
else:
# Note that 0.5 is an empirical value here because we don't know
# the actual number of tokens generated before EOS.
return request_length + 0.5 * max_tokens
proxy_state = None
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument("--port", type=int, default=8000)
parser.add_argument("--host", type=str, default="localhost")
parser.add_argument("--dp-hosts", type=str, nargs="+", default=["localhost"])
parser.add_argument("--dp-ports", type=int, nargs="+", default=[8001])
parser.add_argument("--max-retries", type=int, default=3, help="Maximum number of retries for HTTP requests")
parser.add_argument(
"--retry-delay", type=float, default=0.001, help="Base delay (seconds) for exponential backoff retries"
)
args = parser.parse_args()
if len(args.dp_hosts) != len(args.dp_ports):
raise ValueError("Number of dp hosts must match number of dp ports")
args.server_instances = list(zip(args.dp_hosts, args.dp_ports))
return args
@asynccontextmanager
async def lifespan(app: FastAPI):
global proxy_state
proxy_state = ProxyState(global_args.server_instances)
print(f"Initialized {len(proxy_state.dp_servers)} dp server clients.")
yield
for p in proxy_state.dp_servers:
await p.client.aclose()
async def listen_for_disconnect(request: Request) -> None:
"""Return if a disconnect message is received"""
while True:
message = await request.receive()
if message["type"] == "http.disconnect":
break
def with_cancellation(handler_func):
@functools.wraps(handler_func)
async def wrapper(*args, **kwargs):
request = kwargs["request"]
handler_task = asyncio.create_task(handler_func(*args, **kwargs))
cancellation_task = asyncio.create_task(listen_for_disconnect(request))
done, pending = await asyncio.wait([handler_task, cancellation_task], return_when=asyncio.FIRST_COMPLETED)
for task in pending:
task.cancel()
if handler_task in done:
return handler_task.result()
return None
return wrapper
app = FastAPI(lifespan=lifespan)
async def stream_service_response_with_retry(
client: httpx.AsyncClient,
endpoint: str,
req_data: dict,
request_id: str,
max_retries: int = 3,
base_delay: float = 0.2,
):
headers = {"Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", "X-Request-Id": request_id}
for attempt in range(1, max_retries + 1):
try:
async with client.stream("POST", endpoint, json=req_data, headers=headers) as response:
response.raise_for_status()
first_chunk_sent = False
async for chunk in response.aiter_bytes():
first_chunk_sent = True
yield chunk
return # Success, exit after streaming
except (httpx.RequestError, httpx.HTTPStatusError) as e:
if attempt < max_retries:
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, e)
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
raise e
except Exception as e:
# If any chunk has been sent, do not retry, just log and drop
if "first_chunk_sent" in locals() and first_chunk_sent:
logger.error("Streaming to client interrupted after response started: %s", e)
return
else:
if attempt < max_retries:
logger.warning("Attempt %s failed for streaming %s: %s", attempt, endpoint, e)
await asyncio.sleep(base_delay * (2 ** (attempt - 1)))
else:
logger.error("All %s attempts failed for streaming %s.", max_retries, endpoint)
raise e
async def _select_instance(api: str, req_data: Any, request_length: int):
# refer to vLLM sampling_params: max_token default value
max_tokens = req_data.get("max_tokens", 16)
ignore_eos = req_data.get("ignore_eos", False)
priority_score = proxy_state.calculate_request_score(request_length, max_tokens=max_tokens, ignore_eos=ignore_eos)
logger.debug(
"Request length: %s, max tokens: %s, ignore_eos: %s, Priority score: %s",
request_length,
max_tokens,
ignore_eos,
priority_score,
)
request_id = await proxy_state.next_req_id()
# Select dp server based on priority score
server_idx = proxy_state.select_server(priority_score)
chosen_server = proxy_state.dp_servers[server_idx]
logger.debug("Choose server %s to process request %s", chosen_server.url, request_id)
return InstanceInfo(
request_id=request_id, server_idx=server_idx, priority_score=priority_score, server_state=chosen_server
)
@dataclass
class InstanceInfo:
request_id: str
server_idx: int
priority_score: float
server_state: ServerState
async def _handle_completions(api: str, request: Request):
try:
req_data = await request.json()
req_body = await request.body()
request_length = len(req_body)
instance_info = await _select_instance(api, req_data, request_length)
async def generate_stream():
nonlocal instance_info
# Only one await per chunk, minimal logic in loop
try:
async for chunk in stream_service_response_with_retry(
instance_info.server_state.client,
api,
req_data,
request_id=instance_info.request_id,
max_retries=global_args.max_retries,
base_delay=global_args.retry_delay,
):
yield chunk
except Exception as e:
logger.error(
"Error during streaming from server %s: %s, the aborted request is: %s.",
instance_info.server_state.url,
e,
instance_info.request_id,
)
# After streaming done, release tokens
proxy_state.release_server(instance_info.server_idx, instance_info.priority_score)
return StreamingResponse(generate_stream(), media_type="application/json")
except Exception as e:
import traceback
exc_info = sys.exc_info()
print(f"Error occurred in external dp proxy server - {api} endpoint")
print(e)
print("".join(traceback.format_exception(*exc_info)))
raise
@app.post("/v1/completions")
@with_cancellation
async def handle_completions(request: Request):
return await _handle_completions("/completions", request)
@app.post("/v1/chat/completions")
@with_cancellation
async def handle_chat_completions(request: Request):
return await _handle_completions("/chat/completions", request)
@app.get("/healthcheck")
async def healthcheck():
return {
"status": "ok",
"dp_instances": len(proxy_state.dp_servers),
}
if __name__ == "__main__":
global global_args
global_args = parse_args()
import uvicorn
uvicorn.run(app, host=global_args.host, port=global_args.port)

65
examples/external_online_dp/launch_online_dp.py
View File

@@ -4,52 +4,19 @@ import os
import subprocess
import sys
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument(
"--dp-size",
type=int,
required=True,
help="Data parallel size."
)
parser.add_argument(
"--tp-size",
type=int,
default=1,
help="Tensor parallel size."
)
parser.add_argument(
"--dp-size-local",
type=int,
default=-1,
help="Local data parallel size."
)
parser.add_argument(
"--dp-rank-start",
type=int,
default=0,
help="Starting rank for data parallel."
)
parser.add_argument(
"--dp-address",
type=str,
required=True,
help="IP address for data parallel master node."
)
parser.add_argument(
"--dp-rpc-port",
type=str,
default=12345,
help="Port for data parallel master node."
)
parser.add_argument(
"--vllm-start-port",
type=int,
default=9000,
help="Starting port for the engine."
)
parser.add_argument("--dp-size", type=int, required=True, help="Data parallel size.")
parser.add_argument("--tp-size", type=int, default=1, help="Tensor parallel size.")
parser.add_argument("--dp-size-local", type=int, default=-1, help="Local data parallel size.")
parser.add_argument("--dp-rank-start", type=int, default=0, help="Starting rank for data parallel.")
parser.add_argument("--dp-address", type=str, required=True, help="IP address for data parallel master node.")
parser.add_argument("--dp-rpc-port", type=str, default=12345, help="Port for data parallel master node.")
parser.add_argument("--vllm-start-port", type=int, default=9000, help="Starting port for the engine.")
return parser.parse_args()
args = parse_args()
dp_size = args.dp_size
tp_size = args.tp_size
@@ -61,11 +28,12 @@ dp_address = args.dp_address
dp_rpc_port = args.dp_rpc_port
vllm_start_port = args.vllm_start_port
def run_command(visiable_devices, dp_rank, vllm_engine_port):
def run_command(visible_devices, dp_rank, vllm_engine_port):
command = [
"bash",
"./run_dp_template.sh",
visiable_devices,
visible_devices,
str(vllm_engine_port),
str(dp_size),
str(dp_rank),
@@ -75,6 +43,7 @@ def run_command(visiable_devices, dp_rank, vllm_engine_port):
]
subprocess.run(command, check=True)
if __name__ == "__main__":
template_path = "./run_dp_template.sh"
if not os.path.exists(template_path):
@@ -86,12 +55,10 @@ if __name__ == "__main__":
for i in range(dp_size_local):
dp_rank = dp_rank_start + i
vllm_engine_port = vllm_start_port + i
visiable_devices = ",".join(str(x) for x in range(i * tp_size, (i + 1) * tp_size))
process = multiprocessing.Process(target=run_command,
args=(visiable_devices, dp_rank,
vllm_engine_port))
visible_devices = ",".join(str(x) for x in range(i * tp_size, (i + 1) * tp_size))
process = multiprocessing.Process(target=run_command, args=(visible_devices, dp_rank, vllm_engine_port))
processes.append(process)
process.start()
for process in processes:
process.join()
process.join()

22
examples/external_online_dp/run_dp_template.sh
View File

@@ -2,7 +2,6 @@ export HCCL_IF_IP=your_ip_here
export GLOO_SOCKET_IFNAME=your_socket_ifname_here
export TP_SOCKET_IFNAME=your_socket_ifname_here
export HCCL_SOCKET_IFNAME=your_socket_ifname_here
export DISAGGREGATED_PREFILL_RANK_TABLE_PATH=your_rank_table_path_here
export VLLM_LOGGING_LEVEL="info"
export OMP_PROC_BIND=false
export OMP_NUM_THREADS=10
@@ -11,8 +10,6 @@ export HCCL_DETERMINISTIC=True
export HCCL_BUFFSIZE=1024
export TASK_QUEUE_ENABLE=1
export VLLM_USE_V1=1
export ASCEND_RT_VISIBLE_DEVICES=$1
vllm serve model_path \
@@ -26,21 +23,10 @@ vllm serve model_path \
--enable-expert-parallel \
--seed 1024 \
--served-model-name dsv3 \
--max-model-len 3500 \
--max-num-batched-tokens 3500 \
--max-num-seqs 28 \
--max-model-len 8192 \
--max-num-batched-tokens 2048 \
--max-num-seqs 16 \
--trust-remote-code \
--gpu-memory-utilization 0.9 \
--quantization ascend \
--speculative-config '{"num_speculative_tokens": 1, "method":"deepseek_mtp"}' \
--kv-transfer-config \
'{"kv_connector": "LLMDataDistCMgrConnector",
"kv_buffer_device": "npu",
"kv_role": "kv_consumer",
"kv_parallel_size": "1",
"kv_port": "20001",
"engine_id": "0",
"kv_connector_module_path": "vllm_ascend.distributed.llmdatadist_c_mgr_connector"
}' \
--additional-config \
'{"ascend_scheduler_config": {"enabled": true}, "torchair_graph_config":{"enabled":true,"enable_kv_nz":false, "graph_batch_size":[28]}, "enable_weight_nz_layout":true, "enable_multistream_moe":false}'
--speculative-config '{"num_speculative_tokens": 1, "method":"mtp"}'

86
examples/offline_data_parallel.py
View File

@@ -61,13 +61,13 @@ from time import sleep
import torch
from vllm import LLM, SamplingParams
from vllm.distributed.parallel_state import ( # noqa E402
destroy_distributed_environment, destroy_model_parallel)
from vllm.utils import get_open_port
from vllm.distributed.parallel_state import destroy_distributed_environment, destroy_model_parallel # noqa E402
from vllm.utils.network_utils import get_open_port
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def parse_args():
import argparse
@@ -78,39 +78,18 @@ def parse_args():
default="ibm-research/PowerMoE-3b",
help="Model name or path",
)
parser.add_argument("--dp-size",
type=int,
default=2,
help="Data parallel size")
parser.add_argument("--tp-size",
type=int,
default=1,
help="Tensor parallel size")
parser.add_argument("--node-size",
type=int,
default=1,
help="Total number of nodes")
parser.add_argument("--node-rank",
type=int,
default=0,
help="Rank of the current node")
parser.add_argument("--master-addr",
type=str,
default="",
help="Master node IP address")
parser.add_argument("--master-port",
type=int,
default=0,
help="Master node port")
parser.add_argument("--enforce-eager",
action="store_true",
help="Enforce eager mode execution.")
parser.add_argument("--trust-remote-code",
action="store_true",
help="Trust remote code.")
parser.add_argument("--enable-expert-parallel",
action="store_true",
help="Enable expert parallel, used in MOE models.")
parser.add_argument("--dp-size", type=int, default=2, help="Data parallel size")
parser.add_argument("--tp-size", type=int, default=1, help="Tensor parallel size")
parser.add_argument("--node-size", type=int, default=1, help="Total number of nodes")
parser.add_argument("--node-rank", type=int, default=0, help="Rank of the current node")
parser.add_argument("--master-addr", type=str, default="", help="Master node IP address")
parser.add_argument("--master-port", type=int, default=0, help="Master node port")
parser.add_argument("--enforce-eager", action="store_true", help="Enforce eager mode execution.")
parser.add_argument("--trust-remote-code", action="store_true", help="Trust remote code.")
parser.add_argument(
"--enable-expert-parallel", action="store_true", help="Enable expert parallel, used in MOE models."
)
parser.add_argument("--quantization", type=str, default="", help="Use quantization models")
return parser.parse_args()
@@ -123,6 +102,7 @@ def cleanup_env_and_memory():
torch.npu.empty_cache()
torch.npu.reset_peak_memory_stats()
def main(
model,
dp_size,
@@ -134,6 +114,7 @@ def main(
enable_expert_parallel,
enforce_eager,
trust_remote_code,
quantization,
):
# DP only support on V1 engine
os.environ["VLLM_DP_RANK"] = str(global_dp_rank)
@@ -142,8 +123,13 @@ def main(
os.environ["VLLM_DP_MASTER_IP"] = dp_master_ip
os.environ["VLLM_DP_MASTER_PORT"] = str(dp_master_port)
# CUDA_VISIBLE_DEVICES for each DP rank is set automatically inside the
# engine processes.
from vllm_ascend.utils import vllm_version_is
_dp_device_ids = None
if not vllm_version_is("0.23.0"):
import torch
_dp_device_ids = [str(i) for i in range(torch.npu.device_count())]
# Sample prompts.
prompts = [
@@ -163,7 +149,7 @@ def main(
def start(rank):
return rank * floor + min(rank, remainder)
prompts = prompts[start(global_dp_rank):start(global_dp_rank + 1)]
prompts = prompts[start(global_dp_rank) : start(global_dp_rank + 1)]
if len(prompts) == 0:
# if any rank has no prompts to process,
# we need to set a placeholder prompt
@@ -174,9 +160,7 @@ def main(
# since we are doing data parallel, every rank can have different
# sampling params. here we set different max_tokens for different
# ranks for demonstration.
sampling_params = SamplingParams(temperature=0.8,
top_p=0.95,
max_tokens=[16, 20][global_dp_rank % 2])
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=[16, 20][global_dp_rank % 2])
# Create an LLM.
llm = LLM(
@@ -185,6 +169,8 @@ def main(
enforce_eager=enforce_eager,
enable_expert_parallel=enable_expert_parallel,
trust_remote_code=trust_remote_code,
quantization=quantization,
**({} if _dp_device_ids is None else {"device_ids": _dp_device_ids}),
)
outputs = llm.generate(prompts, sampling_params)
# Print the outputs.
@@ -194,14 +180,14 @@ def main(
break
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"DP rank {global_dp_rank}, Prompt: {prompt!r}, "
f"Generated text: {generated_text!r}")
print(f"DP rank {global_dp_rank}, Prompt: {prompt!r}, Generated text: {generated_text!r}")
# Give engines time to pause their processing loops before exiting.
sleep(5)
del llm
cleanup_env_and_memory()
if __name__ == "__main__":
args = parse_args()
@@ -220,11 +206,12 @@ if __name__ == "__main__":
assert dp_size % node_size == 0, "dp_size should be divisible by node_size"
dp_per_node = dp_size // node_size
quantization = args.quantization if args.quantization else None
from multiprocessing import Process
procs = []
for local_dp_rank, global_dp_rank in enumerate(
range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)):
for local_dp_rank, global_dp_rank in enumerate(range(node_rank * dp_per_node, (node_rank + 1) * dp_per_node)):
proc = Process(
target=main,
args=(
@@ -238,17 +225,16 @@ if __name__ == "__main__":
args.enable_expert_parallel,
args.enforce_eager,
args.trust_remote_code,
quantization,
),
)
proc.start()
procs.append(proc)
exit_code = 0
for proc in procs:
proc.join(timeout=300)
proc.join(timeout=900)
if proc.exitcode is None:
print(
f"Killing process {proc.pid} that didn't stop within 5 minutes."
)
print(f"Killing process {proc.pid} that didn't stop within 15 minutes.")
proc.kill()
exit_code = 1
elif proc.exitcode:

83
examples/offline_disaggregated_prefill_npu.py
View File

@@ -24,12 +24,13 @@ from multiprocessing import Event, Process
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def clean_up():
import gc
import torch
from vllm.distributed.parallel_state import (
destroy_distributed_environment, destroy_model_parallel)
from vllm.distributed.parallel_state import destroy_distributed_environment, destroy_model_parallel
destroy_model_parallel()
destroy_distributed_environment()
gc.collect()
@@ -37,29 +38,34 @@ def clean_up():
def run_prefill(prefill_done, process_close):
# ranktable.json needs be generated using gen_ranktable.sh
# from the examples/disaggregated_prefill_v1 in the main branch.
os.environ['DISAGGREGATED_PREFILL_RANK_TABLE_PATH'] = "./ranktable.json"
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "0"
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
prompts = [
"Hello, how are you today?", "Hi, what is your name?",
"Tell me a very long story.", "what is your favourite book?"
"Hello, how are you today?",
"Hi, what is your name?",
"Tell me a very long story.",
"what is your favourite book?",
]
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=1)
ktc = KVTransferConfig(kv_connector="LLMDataDistCMgrConnector", kv_buffer_device="npu", kv_role="kv_producer",
kv_parallel_size=1,
kv_connector_module_path="vllm_ascend.distributed.llmdatadist_c_mgr_connector")
ktc = KVTransferConfig(
kv_connector="MooncakeConnectorV1",
kv_role="kv_producer",
kv_port="30000",
engine_id="0",
kv_connector_extra_config={"prefill": {"dp_size": 1, "tp_size": 1}, "decode": {"dp_size": 1, "tp_size": 1}},
)
# Set NPU memory utilization to 0.8
llm = LLM(model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
kv_transfer_config=ktc,
max_model_len=2000,
gpu_memory_utilization=0.8,
tensor_parallel_size=1)
llm = LLM(
model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
kv_transfer_config=ktc,
max_model_len=2000,
gpu_memory_utilization=0.8,
tensor_parallel_size=1,
)
llm.generate(prompts, sampling_params)
print("Prefill node is finished.")
@@ -79,29 +85,34 @@ def run_prefill(prefill_done, process_close):
def run_decode(prefill_done):
os.environ['VLLM_ASCEND_LLMDD_RPC_PORT'] = '6634'
# ranktable.json needs be generated using gen_ranktable.sh
# from the examples/disaggregated_prefill_v1 module in the main branch.
os.environ['DISAGGREGATED_PREFILL_RANK_TABLE_PATH'] = "./ranktable.json"
os.environ["ASCEND_RT_VISIBLE_DEVICES"] = "1"
from vllm import LLM, SamplingParams
from vllm.config import KVTransferConfig
prompts = [
"Hello, how are you today?", "Hi, what is your name?",
"Tell me a very long story.", "what is your favourite book?"
"Hello, how are you today?",
"Hi, what is your name?",
"Tell me a very long story.",
"what is your favourite book?",
]
sampling_params = SamplingParams(temperature=0, top_p=0.95)
ktc = KVTransferConfig(kv_connector="LLMDataDistCMgrConnector", kv_buffer_device="npu", kv_role="kv_consumer",
kv_parallel_size=1, kv_connector_module_path="vllm_ascend.distributed.llmdatadist_c_mgr_connector")
ktc = KVTransferConfig(
kv_connector="MooncakeConnectorV1",
kv_role="kv_consumer",
kv_port="30100",
engine_id="1",
kv_connector_extra_config={"prefill": {"dp_size": 1, "tp_size": 1}, "decode": {"dp_size": 1, "tp_size": 1}},
)
llm = LLM(model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
kv_transfer_config=ktc,
max_model_len=2000,
gpu_memory_utilization=0.8,
tensor_parallel_size=1)
llm = LLM(
model="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
kv_transfer_config=ktc,
max_model_len=2000,
gpu_memory_utilization=0.8,
tensor_parallel_size=1,
)
# Wait for the producer to start the consumer
print("Waiting for prefill node to finish...")
@@ -120,16 +131,18 @@ def run_decode(prefill_done):
if __name__ == "__main__":
mp.get_context('spawn')
mp.get_context("spawn")
prefill_done = Event()
process_close = Event()
prefill_process = Process(target=run_prefill,
args=(
prefill_done,
process_close,
))
decode_process = Process(target=run_decode, args=(prefill_done, ))
prefill_process = Process(
target=run_prefill,
args=(
prefill_done,
process_close,
),
)
decode_process = Process(target=run_decode, args=(prefill_done,))
# Start prefill node
prefill_process.start()

16
examples/offline_embed.py
View File

@@ -25,31 +25,33 @@ from vllm import LLM
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def get_detailed_instruct(task_description: str, query: str) -> str:
return f'Instruct: {task_description}\nQuery:{query}'
return f"Instruct: {task_description}\nQuery:{query}"
def main():
# Each query must come with a one-sentence instruction that describes the task
task = 'Given a web search query, retrieve relevant passages that answer the query'
task = "Given a web search query, retrieve relevant passages that answer the query"
queries = [
get_detailed_instruct(task, 'What is the capital of China?'),
get_detailed_instruct(task, 'Explain gravity')
get_detailed_instruct(task, "What is the capital of China?"),
get_detailed_instruct(task, "Explain gravity"),
]
# No need to add instruction for retrieval documents
documents = [
"The capital of China is Beijing.",
"Gravity is a force that attracts two bodies towards each other. It gives weight to physical objects and is responsible for the movement of planets around the sun."
"Gravity is a force that attracts two bodies towards each other. "
"It gives weight to physical objects and is responsible for the movement of planets around the sun.",
]
input_texts = queries + documents
model = LLM(model="Qwen/Qwen3-Embedding-0.6B", task="embed")
model = LLM(model="Qwen/Qwen3-Embedding-0.6B", runner="pooling")
outputs = model.embed(input_texts)
embeddings = torch.tensor([o.outputs.embedding for o in outputs])
# Calculate the similarity scores between the first two queries and the last two documents
scores = (embeddings[:2] @ embeddings[2:].T)
scores = embeddings[:2] @ embeddings[2:].T
print(scores.tolist())
# [[0.7620252966880798, 0.14078938961029053], [0.1358368694782257, 0.6013815999031067]]

175
examples/offline_external_launcher.py
View File

@@ -63,17 +63,47 @@ from multiprocessing import Process
from time import sleep
import torch
from safetensors.torch import load_file
from vllm import LLM, SamplingParams
from vllm.distributed.parallel_state import ( # noqa E402
destroy_distributed_environment, destroy_model_parallel, get_tp_group)
from vllm.utils import get_open_port, GiB_bytes
destroy_distributed_environment,
destroy_model_parallel,
get_tp_group,
)
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.network_utils import get_open_port
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def parse_args():
def patch_vllm_moe_model_weight_loader(model):
model = getattr(model, "model", None) or getattr(model, "language_model", None)
if model is None:
raise ValueError("The provided model does not have a valid 'model' or 'language_model' attribute.")
for layer in model.layers:
mlp_attr = "mlp"
mlp = getattr(layer, mlp_attr)
param_dict = dict(mlp.named_parameters())
for name, param in param_dict.items():
if "w13_weight" in name or "w2_weight" in name:
param.weight_loader = mlp.experts.weight_loader
def load_and_merge_safetensors(directory):
if not os.path.isdir(directory):
raise ValueError(f"The provided directory does not exist: {directory}")
merged_dict = {}
for filename in os.listdir(directory):
if filename.endswith(".safetensors"):
file_path = os.path.join(directory, filename)
print(f"loading file: {file_path}")
f = load_file(file_path)
merged_dict.update(f)
return merged_dict
def parse_args():
parser = argparse.ArgumentParser(description="External launcher Inference")
parser.add_argument(
"--model",
@@ -81,55 +111,41 @@ def parse_args():
default="Qwen/Qwen3-0.6B",
help="Model name or path",
)
parser.add_argument("--tp-size",
type=int,
default=1,
help="Tensor parallel size")
parser.add_argument("--node-size",
type=int,
default=1,
help="Total number of nodes")
parser.add_argument("--node-rank",
type=int,
default=0,
help="Rank of the current node")
parser.add_argument("--proc-per-node",
type=int,
default=1,
help="Number of processes per node")
parser.add_argument("--master-addr",
type=str,
default="",
help="Master node IP address")
parser.add_argument("--master-port",
type=int,
default=0,
help="Master node port")
parser.add_argument("--enforce-eager",
action="store_true",
help="Enforce eager mode execution.")
parser.add_argument("--trust-remote-code",
action="store_true",
help="Trust remote code.")
parser.add_argument("--enable-expert-parallel",
action="store_true",
help="Enable expert parallel, used in MOE models.")
parser.add_argument("--enable-sleep-mode",
action="store_true",
help="Enable sleep mode for the engine.")
parser.add_argument("--temperature",
type=float,
default=0.8,
help="Float that controls the randomness of the sampling.")
parser.add_argument("--model-weight-gib",
type=float,
default=None,
help="Model weight memory usage in GiB (e.g., 1.0 for 0.5B model).")
parser.add_argument("--tp-size", type=int, default=1, help="Tensor parallel size")
parser.add_argument("--node-size", type=int, default=1, help="Total number of nodes")
parser.add_argument("--node-rank", type=int, default=0, help="Rank of the current node")
parser.add_argument("--proc-per-node", type=int, default=1, help="Number of processes per node")
parser.add_argument("--master-addr", type=str, default="", help="Master node IP address")
parser.add_argument("--master-port", type=int, default=0, help="Master node port")
parser.add_argument("--enforce-eager", action="store_true", help="Enforce eager mode execution.")
parser.add_argument("--trust-remote-code", action="store_true", help="Trust remote code.")
parser.add_argument(
"--enable-expert-parallel", action="store_true", help="Enable expert parallel, used in MOE models."
)
parser.add_argument("--enable-sleep-mode", action="store_true", help="Enable sleep mode for the engine.")
parser.add_argument(
"--temperature", type=float, default=0.8, help="Float that controls the randomness of the sampling."
)
parser.add_argument(
"--model-weight-gib",
type=float,
default=None,
help="Model weight memory usage in GiB (e.g., 1.0 for 0.5B model).",
)
parser.add_argument(
"--sleep-mode-level",
type=int,
choices=[1, 2],
default=1,
help="Sleep mode level: 1 or 2. This example of level 2 is only supported for dense model.",
)
args = parser.parse_args()
if args.enable_sleep_mode:
if args.model_weight_gib is None or args.temperature != 0:
parser.error("model-weight-gib must be provided, and temperature must be zero when enable-sleep-mode is set.")
parser.error(
"model-weight-gib must be provided, and temperature must be zero when enable-sleep-mode is set."
)
if args.model_weight_gib <= 0:
parser.error("model-weight-gib must be greater than 0 when enable-sleep-mode is set.")
if args.model == parser.get_default("model") and args.model_weight_gib is None:
@@ -152,6 +168,7 @@ def main(
trust_remote_code: bool = True,
enable_sleep_mode: bool = False,
temperature: float = 0.8,
sleep_mode_level: int = 1,
):
os.environ["MASTER_ADDR"] = master_addr
os.environ["MASTER_PORT"] = str(master_port)
@@ -186,22 +203,31 @@ def main(
enable_sleep_mode=enable_sleep_mode,
)
tp_ranks = get_tp_group().ranks
print(f'TP RANKS: {tp_ranks}')
print(f"TP RANKS: {tp_ranks}")
outputs = llm.generate(prompts, sampling_params)
if enable_sleep_mode:
if rank == 0:
free_bytes_before_sleep, total = torch.npu.mem_get_info()
llm.sleep(level=1)
llm.sleep(level=sleep_mode_level)
if rank == 0:
free_bytes_after_sleep, total = torch.npu.mem_get_info()
freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
print(f"Freed memory: {freed_bytes / 1024 ** 3:.2f} GiB")
print(f"Freed memory: {freed_bytes / 1024**3:.2f} GiB")
# now the freed memory should be larger than the model weights
assert freed_bytes >= model_weight_gib / tensor_parallel_size * GiB_bytes
llm.wake_up()
if sleep_mode_level == 1:
llm.wake_up()
else:
llm.wake_up(tags=["weights"])
run_model = llm.llm_engine.model_executor.driver_worker.worker.model_runner.model
patch_vllm_moe_model_weight_loader(run_model)
sd = load_and_merge_safetensors(model)
run_model.load_weights(sd.items())
llm.wake_up(tags=["kv_cache"])
outputs_after_wakeup = llm.generate(prompts, sampling_params)
if rank == 0:
# cmp output
@@ -214,8 +240,7 @@ def main(
break
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Global rank: {rank}, Prompt: {prompt!r}, "
f"Generated text: {generated_text!r}")
print(f"Global rank: {rank}, Prompt: {prompt!r}, Generated text: {generated_text!r}")
# Give engines time to pause their processing loops before exiting.
sleep(5)
@@ -251,24 +276,26 @@ if __name__ == "__main__":
world_size = node_size * proc_per_node
procs = []
for local_rank, rank in enumerate(
range(proc_per_node * node_rank, proc_per_node * (node_rank + 1))):
proc = Process(target=main,
args=(
local_rank,
rank,
master_addr,
master_port,
args.model_weight_gib,
args.model,
world_size,
tp_size,
args.enable_expert_parallel,
args.enforce_eager,
args.trust_remote_code,
args.enable_sleep_mode,
args.temperature,
))
for local_rank, rank in enumerate(range(proc_per_node * node_rank, proc_per_node * (node_rank + 1))):
proc = Process(
target=main,
args=(
local_rank,
rank,
master_addr,
master_port,
args.model_weight_gib,
args.model,
world_size,
tp_size,
args.enable_expert_parallel,
args.enforce_eager,
args.trust_remote_code,
args.enable_sleep_mode,
args.temperature,
args.sleep_mode_level,
),
)
proc.start()
procs.append(proc)
@@ -276,9 +303,7 @@ if __name__ == "__main__":
for proc in procs:
proc.join(timeout=600)
if proc.exitcode is None:
print(
f"Killing process {proc.pid} that didn't stop within 30 minutes."
)
print(f"Killing process {proc.pid} that didn't stop within 30 minutes.")
proc.kill()
exit_code = 1
elif proc.exitcode:

51
examples/offline_inference_audio_language.py
View File

@@ -17,19 +17,20 @@
# Adapted from vllm-project/vllm/examples/offline_inference/audio_language.py
#
"""
This example shows how to use vLLM for running offline inference
This example shows how to use vLLM for running offline inference
with the correct prompt format on audio language models.
For most models, the prompt format should follow corresponding examples
on HuggingFace model repository.
"""
import os
import argparse
import os
from vllm.assets.audio import AudioAsset
try:
import librosa # type: ignore
import librosa # type: ignore
except ImportError:
raise Exception("Can't import librosa, please ensure it's installed")
@@ -40,7 +41,7 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def prepare_inputs(audio_count: int, audio_path1: str, audio_path2: str):
use_vllm_audio_assert = True if audio_path1 == "mary_had_lamb" and audio_path2 == "winning_call" else False
use_vllm_audio_assert = audio_path1 == "mary_had_lamb" and audio_path2 == "winning_call"
if use_vllm_audio_assert:
audio_assets = [AudioAsset("mary_had_lamb"), AudioAsset("winning_call")]
else:
@@ -48,22 +49,22 @@ def prepare_inputs(audio_count: int, audio_path1: str, audio_path2: str):
question_per_audio_count = {
1: "What is recited in the audio?",
2: "What sport and what nursery rhyme are referenced?"
2: "What sport and what nursery rhyme are referenced?",
}
audio_in_prompt = "".join([
f"Audio {idx+1}: <|audio_bos|><|AUDIO|><|audio_eos|>\n"
for idx in range(audio_count)
])
audio_in_prompt = "".join([f"Audio {idx + 1}: <|audio_bos|><|AUDIO|><|audio_eos|>\n" for idx in range(audio_count)])
question = question_per_audio_count[audio_count]
prompt = ("<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
"<|im_start|>user\n"
f"{audio_in_prompt}{question}<|im_end|>\n"
"<|im_start|>assistant\n")
prompt = (
"<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n"
"<|im_start|>user\n"
f"{audio_in_prompt}{question}<|im_end|>\n"
"<|im_start|>assistant\n"
)
mm_data = {
"audio":
audio_assets if not use_vllm_audio_assert else [asset.audio_and_sample_rate for asset in audio_assets[:audio_count]]
"audio": audio_assets
if not use_vllm_audio_assert
else [asset.audio_and_sample_rate for asset in audio_assets[:audio_count]]
}
# Merge text prompt and audio data into inputs
@@ -76,17 +77,17 @@ def main(audio_count: int, audio_path1: str, audio_path2: str):
# lower-end GPUs.
# Unless specified, these settings have been tested to work on a single L4.
# `limit_mm_per_prompt`: the max num items for each modality per prompt.
llm = LLM(model="Qwen/Qwen2-Audio-7B-Instruct",
max_model_len=4096,
max_num_seqs=5,
limit_mm_per_prompt={"audio": audio_count},
enforce_eager=True)
llm = LLM(
model="Qwen/Qwen2-Audio-7B-Instruct",
max_model_len=4096,
max_num_seqs=5,
limit_mm_per_prompt={"audio": audio_count},
enforce_eager=True,
)
inputs = prepare_inputs(audio_count, audio_path1, audio_path2)
sampling_params = SamplingParams(temperature=0.2,
max_tokens=64,
stop_token_ids=None)
sampling_params = SamplingParams(temperature=0.2, max_tokens=64, stop_token_ids=None)
outputs = llm.generate(inputs, sampling_params=sampling_params)
@@ -96,7 +97,9 @@ def main(audio_count: int, audio_path1: str, audio_path2: str):
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Arguments of rank table generator", )
parser = argparse.ArgumentParser(
description="Arguments of rank table generator",
)
parser.add_argument("--audio-path1", type=str, default="mary_had_lamb")
parser.add_argument("--audio-path2", type=str, default="winning_call")
args = parser.parse_args()

72
examples/offline_inference_metrics.py Normal file
View File

@@ -0,0 +1,72 @@
#
# Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
Example: Access request-level metrics from vLLM outputs.
By default, vLLM disables log stats (disable_log_stats=True), which causes
output.metrics to be None. To populate metrics such as first_token_time,
finished_time, etc., you must explicitly set disable_log_stats=False when
creating the LLM instance.
See: https://github.com/vllm-project/vllm-ascend/issues/5027
"""
# isort: skip_file
import os
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
from vllm import LLM, SamplingParams
def main():
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Create a sampling params object.
sampling_params = SamplingParams(max_tokens=100, temperature=0.0)
# IMPORTANT: Set disable_log_stats=False to enable output.metrics.
# Without this, output.metrics will be None.
llm = LLM(model="Qwen/Qwen2.5-0.5B-Instruct", disable_log_stats=False)
# Generate texts from the prompts.
outputs = llm.generate(prompts, sampling_params)
for output in outputs:
prompt = output.prompt
generated_text = output.outputs[0].text
metrics = output.metrics
print(f"Prompt: {prompt!r}")
print(f" Generated text: {generated_text!r}")
if metrics is not None:
print(f" Arrival time: {metrics.arrival_time}")
print(f" First scheduled time: {metrics.first_scheduled_time}")
print(f" First token time: {metrics.first_token_time}")
print(f" Finished time: {metrics.finished_time}")
else:
print(" Metrics: None (set disable_log_stats=False to enable)")
print()
if __name__ == "__main__":
main()

3
examples/offline_inference_npu.py
View File

@@ -37,6 +37,9 @@ def main():
# Create a sampling params object.
sampling_params = SamplingParams(max_tokens=100, temperature=0.0)
# Create an LLM.
# NOTE: To access output.metrics (e.g., first_token_time, finished_time),
# set disable_log_stats=False. By default, vLLM disables log stats and
# output.metrics will be None. See issue #5027 for details.
llm = LLM(model="Qwen/Qwen2.5-0.5B-Instruct")
# Generate texts from the prompts.

58
examples/offline_inference_npu_long_seq.py Normal file
View File

@@ -0,0 +1,58 @@
import argparse
import os
import time
from vllm import LLM, SamplingParams
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument("--input_len", type=int, default=1024)
parser.add_argument("--output_len", type=int, default=128)
parser.add_argument("--bs", type=int, default=1)
parser.add_argument("--model_path", type=str, default="deepseek-ai/DeepSeek-V2-Lite")
parser.add_argument("--tp", type=int, default=2)
parser.add_argument("--pcp", type=int, default=2)
parser.add_argument("--dcp", type=int, default=1)
parser.add_argument("--iter_times", type=int, default=1)
args = parser.parse_args()
prompts = [
"The capital of France is",
"Hello, my name is Tom, I am",
"The president of United States is",
"AI future is",
]
sampling_params = SamplingParams(temperature=0.8, top_p=0.95, max_tokens=args.output_len)
llm = LLM(
model=args.model_path,
trust_remote_code=True,
enforce_eager=True,
tensor_parallel_size=args.tp,
prefill_context_parallel_size=args.pcp,
decode_context_parallel_size=args.dcp,
enable_prefix_caching=False,
enable_expert_parallel=True,
enable_chunked_prefill=False,
max_num_batched_tokens=2048,
max_model_len=1024,
max_num_seqs=1,
block_size=128,
gpu_memory_utilization=0.9,
)
t0 = time.time()
for _ in range(args.iter_times):
outputs = llm.generate(prompts, sampling_params)
t1 = time.time()
print(f"TTFT: {(t1 - t0) * 1000 / (args.iter_times * args.bs)} ms")
for i, output in enumerate(outputs):
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"req_num: {i}\nGenerated text: {generated_text!r}")

12
examples/offline_inference_npu_tp2.py
View File

@@ -37,11 +37,13 @@ def main():
# Create a sampling params object.
sampling_params = SamplingParams(max_tokens=100, temperature=0.0)
# Create an LLM.
llm = LLM(model="deepseek-ai/DeepSeek-V2-Lite",
tensor_parallel_size=2,
enforce_eager=True,
trust_remote_code=True,
max_model_len=1024)
llm = LLM(
model="deepseek-ai/DeepSeek-V2-Lite",
tensor_parallel_size=2,
enforce_eager=True,
trust_remote_code=True,
max_model_len=1024,
)
# Generate texts from the prompts.
outputs = llm.generate(prompts, sampling_params)

9
examples/offline_inference_sleep_mode_npu.py
View File

@@ -20,16 +20,17 @@ import os
import torch
from vllm import LLM, SamplingParams
from vllm.utils import GiB_bytes
from vllm.utils.mem_constants import GiB_bytes
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def main():
prompt = "How are you?"
free, total = torch.npu.mem_get_info()
print(f"Free memory before sleep: {free / 1024 ** 3:.2f} GiB")
print(f"Free memory before sleep: {free / 1024**3:.2f} GiB")
# record npu memory use baseline in case other process is running
used_bytes_baseline = total - free
llm = LLM("Qwen/Qwen2.5-0.5B-Instruct", enable_sleep_mode=True)
@@ -39,9 +40,7 @@ def main():
llm.sleep(level=1)
free_npu_bytes_after_sleep, total = torch.npu.mem_get_info()
print(
f"Free memory after sleep: {free_npu_bytes_after_sleep / 1024 ** 3:.2f} GiB"
)
print(f"Free memory after sleep: {free_npu_bytes_after_sleep / 1024**3:.2f} GiB")
used_bytes = total - free_npu_bytes_after_sleep - used_bytes_baseline
# now the memory usage should be less than the model weights
# (0.5B model, 1GiB weights)

165
examples/offline_weight_load.py
View File

@@ -63,15 +63,21 @@ from multiprocessing import Process
from time import sleep
import torch
from safetensors.torch import load_file
from vllm import LLM, SamplingParams
from vllm.distributed.parallel_state import ( # noqa E402
destroy_distributed_environment, destroy_model_parallel, get_tp_group)
from vllm.utils import get_open_port, GiB_bytes
from safetensors.torch import load_file
destroy_distributed_environment,
destroy_model_parallel,
get_tp_group,
)
from vllm.model_executor.model_loader.utils import process_weights_after_loading
from vllm.utils.mem_constants import GiB_bytes
from vllm.utils.network_utils import get_open_port
os.environ["VLLM_USE_MODELSCOPE"] = "True"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def patch_vllm_moe_model_weight_loader(model):
# Define MLP attribute mapping for different model types
@@ -88,24 +94,25 @@ def patch_vllm_moe_model_weight_loader(model):
if "w13_weight" in name or "w2_weight" in name:
param.weight_loader = mlp.experts.weight_loader
def load_and_merge_safetensors(directory):
merged_dict = {}
if not os.path.isdir(directory):
raise ValueError(f"directory is not exist : {directory}")
for filename in os.listdir(directory):
if filename.endswith('.safetensors'):
if filename.endswith(".safetensors"):
file_path = os.path.join(directory, filename)
print(f"loading file: {file_path}")
f = load_file(file_path)
merged_dict.update(f)
return merged_dict
def parse_args():
def parse_args():
parser = argparse.ArgumentParser(description="External launcher Inference")
parser.add_argument(
"--model",
@@ -113,55 +120,34 @@ def parse_args():
default="Qwen/Qwen3-0.6B",
help="Model name or path",
)
parser.add_argument("--tp-size",
type=int,
default=1,
help="Tensor parallel size")
parser.add_argument("--node-size",
type=int,
default=1,
help="Total number of nodes")
parser.add_argument("--node-rank",
type=int,
default=0,
help="Rank of the current node")
parser.add_argument("--proc-per-node",
type=int,
default=1,
help="Number of processes per node")
parser.add_argument("--master-addr",
type=str,
default="",
help="Master node IP address")
parser.add_argument("--master-port",
type=int,
default=0,
help="Master node port")
parser.add_argument("--enforce-eager",
action="store_true",
help="Enforce eager mode execution.")
parser.add_argument("--trust-remote-code",
action="store_true",
help="Trust remote code.")
parser.add_argument("--enable-expert-parallel",
action="store_true",
help="Enable expert parallel, used in MOE models.")
parser.add_argument("--enable-sleep-mode",
action="store_true",
help="Enable sleep mode for the engine.")
parser.add_argument("--temperature",
type=float,
default=0.8,
help="Float that controls the randomness of the sampling.")
parser.add_argument("--model-weight-gib",
type=float,
default=None,
help="Model weight memory usage in GiB (e.g., 1.0 for 0.5B model).")
parser.add_argument("--tp-size", type=int, default=1, help="Tensor parallel size")
parser.add_argument("--node-size", type=int, default=1, help="Total number of nodes")
parser.add_argument("--node-rank", type=int, default=0, help="Rank of the current node")
parser.add_argument("--proc-per-node", type=int, default=1, help="Number of processes per node")
parser.add_argument("--master-addr", type=str, default="", help="Master node IP address")
parser.add_argument("--master-port", type=int, default=0, help="Master node port")
parser.add_argument("--enforce-eager", action="store_true", help="Enforce eager mode execution.")
parser.add_argument("--trust-remote-code", action="store_true", help="Trust remote code.")
parser.add_argument(
"--enable-expert-parallel", action="store_true", help="Enable expert parallel, used in MOE models."
)
parser.add_argument("--enable-sleep-mode", action="store_true", help="Enable sleep mode for the engine.")
parser.add_argument(
"--temperature", type=float, default=0.8, help="Float that controls the randomness of the sampling."
)
parser.add_argument(
"--model-weight-gib",
type=float,
default=None,
help="Model weight memory usage in GiB (e.g., 1.0 for 0.5B model).",
)
args = parser.parse_args()
if args.enable_sleep_mode:
if args.model_weight_gib is None or args.temperature != 0:
parser.error("model-weight-gib must be provided, and temperature must be zero when enable-sleep-mode is set.")
parser.error(
"model-weight-gib must be provided, and temperature must be zero when enable-sleep-mode is set."
)
if args.model_weight_gib <= 0:
parser.error("model-weight-gib must be greater than 0 when enable-sleep-mode is set.")
if args.model == parser.get_default("model") and args.model_weight_gib is None:
@@ -215,18 +201,9 @@ def main(
trust_remote_code=trust_remote_code,
distributed_executor_backend="external_launcher",
seed=0,
gpu_memory_utilization = 0.95,
gpu_memory_utilization=0.95,
enable_sleep_mode=enable_sleep_mode,
)
model_path = model
runmodel = llm.llm_engine.model_executor.driver_worker.worker.model_runner.model
patch_vllm_moe_model_weight_loader(runmodel)
sd = load_and_merge_safetensors(model_path)
runmodel.load_weights(sd.items())
print('load state dict done')
tp_ranks = get_tp_group().ranks
print(f'TP RANKS: {tp_ranks}')
outputs = llm.generate(prompts, sampling_params)
if enable_sleep_mode:
@@ -236,11 +213,25 @@ def main(
if rank == 0:
free_bytes_after_sleep, total = torch.npu.mem_get_info()
freed_bytes = free_bytes_after_sleep - free_bytes_before_sleep
print(f"Freed memory: {freed_bytes / 1024 ** 3:.2f} GiB")
print(f"Freed memory: {freed_bytes / 1024**3:.2f} GiB")
# now the freed memory should be larger than the model weights
assert freed_bytes >= model_weight_gib / tensor_parallel_size * GiB_bytes
llm.wake_up()
model_path = model
runmodel = llm.llm_engine.model_executor.driver_worker.worker.model_runner.model
patch_vllm_moe_model_weight_loader(runmodel)
sd = load_and_merge_safetensors(model_path)
runmodel.load_weights(sd.items())
print("load state dict done")
tp_ranks = get_tp_group().ranks
print(f"TP RANKS: {tp_ranks}")
vllm_config = llm.llm_engine.vllm_config.model_config
device = next(runmodel.parameters()).device
process_weights_after_loading(runmodel, vllm_config, device)
outputs_after_wakeup = llm.generate(prompts, sampling_params)
if rank == 0:
# cmp output
@@ -253,8 +244,7 @@ def main(
break
prompt = output.prompt
generated_text = output.outputs[0].text
print(f"Global rank: {rank}, Prompt: {prompt!r}, "
f"Generated text: {generated_text!r}")
print(f"Global rank: {rank}, Prompt: {prompt!r}, Generated text: {generated_text!r}")
# Give engines time to pause their processing loops before exiting.
sleep(5)
@@ -290,24 +280,25 @@ if __name__ == "__main__":
world_size = node_size * proc_per_node
procs = []
for local_rank, rank in enumerate(
range(proc_per_node * node_rank, proc_per_node * (node_rank + 1))):
proc = Process(target=main,
args=(
local_rank,
rank,
master_addr,
master_port,
args.model_weight_gib,
args.model,
world_size,
tp_size,
args.enable_expert_parallel,
args.enforce_eager,
args.trust_remote_code,
args.enable_sleep_mode,
args.temperature,
))
for local_rank, rank in enumerate(range(proc_per_node * node_rank, proc_per_node * (node_rank + 1))):
proc = Process(
target=main,
args=(
local_rank,
rank,
master_addr,
master_port,
args.model_weight_gib,
args.model,
world_size,
tp_size,
args.enable_expert_parallel,
args.enforce_eager,
args.trust_remote_code,
args.enable_sleep_mode,
args.temperature,
),
)
proc.start()
procs.append(proc)
@@ -315,9 +306,7 @@ if __name__ == "__main__":
for proc in procs:
proc.join(timeout=600)
if proc.exitcode is None:
print(
f"Killing process {proc.pid} that didn't stop within 30 minutes."
)
print(f"Killing process {proc.pid} that didn't stop within 30 minutes.")
proc.kill()
exit_code = 1
elif proc.exitcode:

88
examples/prompt_embed_inference.py Normal file
View File

@@ -0,0 +1,88 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates how to generate prompt embeddings using
Hugging Face Transformers and use them as input to vLLM
for both single and batch inference.
Model: meta-llama/Llama-3.2-1B-Instruct
Note: This model is gated on Hugging Face Hub.
You must request access to use it:
https://huggingface.co/meta-llama/Llama-3.2-1B-Instruct
Requirements:
- vLLM
- transformers
Run:
python examples/prompt_embed_inference.py
"""
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer
from vllm import LLM
def init_tokenizer_and_llm(model_name: str):
llm = LLM(model=model_name, enable_prompt_embeds=True)
tokenizer = AutoTokenizer.from_pretrained(model_name)
transformers_model = AutoModelForCausalLM.from_pretrained(model_name)
embedding_layer = transformers_model.get_input_embeddings()
return tokenizer, embedding_layer, llm
def get_prompt_embeds(
chat: list[dict[str, str]],
tokenizer: PreTrainedTokenizer,
embedding_layer: torch.nn.Module,
):
token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors="pt", return_dict=False)
prompt_embeds = embedding_layer(token_ids).squeeze(0)
return prompt_embeds
def single_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module):
chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
prompt_embeds = get_prompt_embeds(chat, tokenizer, embedding_layer)
outputs = llm.generate(
{
"prompt_embeds": prompt_embeds,
}
)
print("\n[Single Inference Output]")
print("-" * 30)
for o in outputs:
print(o.outputs[0].text)
print("-" * 30)
def batch_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module):
chats = [
[{"role": "user", "content": "Please tell me about the capital of France."}],
[{"role": "user", "content": "When is the day longest during the year?"}],
[{"role": "user", "content": "Where is bigger, the moon or the sun?"}],
]
prompt_embeds_list = [get_prompt_embeds(chat, tokenizer, embedding_layer) for chat in chats]
outputs = llm.generate([{"prompt_embeds": embeds} for embeds in prompt_embeds_list])
print("\n[Batch Inference Outputs]")
print("-" * 30)
for i, o in enumerate(outputs):
print(f"Q{i + 1}: {chats[i][0]['content']}")
print(f"A{i + 1}: {o.outputs[0].text}\n")
print("-" * 30)
def main():
model_name = "meta-llama/Llama-3.2-1B-Instruct"
tokenizer, embedding_layer, llm = init_tokenizer_and_llm(model_name)
single_prompt_inference(llm, tokenizer, embedding_layer)
batch_prompt_inference(llm, tokenizer, embedding_layer)
if __name__ == "__main__":
main()

59
examples/prompt_embedding_inference.py
View File

@@ -1,8 +1,7 @@
import os
import torch
from transformers import (AutoModelForCausalLM, AutoTokenizer,
PreTrainedTokenizer)
from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedTokenizer
from vllm import LLM
os.environ["VLLM_USE_MODELSCOPE"] = "True"
@@ -17,27 +16,21 @@ def init_tokenizer_and_llm(model_name: str):
return tokenizer, embedding_layer, llm
def get_prompt_embeds(chat: list[dict[str,
str]], tokenizer: PreTrainedTokenizer,
embedding_layer: torch.nn.Module):
token_ids = tokenizer.apply_chat_template(chat,
add_generation_prompt=True,
return_tensors='pt')
def get_prompt_embeds(chat: list[dict[str, str]], tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module):
token_ids = tokenizer.apply_chat_template(chat, add_generation_prompt=True, return_tensors="pt", return_dict=False)
prompt_embeds = embedding_layer(token_ids).squeeze(0)
return prompt_embeds
def single_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer,
embedding_layer: torch.nn.Module):
chat = [{
"role": "user",
"content": "Please tell me about the capital of France."
}]
def single_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module):
chat = [{"role": "user", "content": "Please tell me about the capital of France."}]
prompt_embeds = get_prompt_embeds(chat, tokenizer, embedding_layer)
outputs = llm.generate({
"prompt_embeds": prompt_embeds,
})
outputs = llm.generate(
{
"prompt_embeds": prompt_embeds,
}
)
print("\n[Single Inference Output]")
print("-" * 30)
@@ -46,34 +39,22 @@ def single_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer,
print("-" * 30)
def batch_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer,
embedding_layer: torch.nn.Module):
chats = [[{
"role": "user",
"content": "Please tell me about the capital of France."
}],
[{
"role": "user",
"content": "When is the day longest during the year?"
}],
[{
"role": "user",
"content": "Where is bigger, the moon or the sun?"
}]]
prompt_embeds_list = [
get_prompt_embeds(chat, tokenizer, embedding_layer) for chat in chats
def batch_prompt_inference(llm: LLM, tokenizer: PreTrainedTokenizer, embedding_layer: torch.nn.Module):
chats = [
[{"role": "user", "content": "Please tell me about the capital of France."}],
[{"role": "user", "content": "When is the day longest during the year?"}],
[{"role": "user", "content": "Where is bigger, the moon or the sun?"}],
]
outputs = llm.generate([{
"prompt_embeds": embeds
} for embeds in prompt_embeds_list])
prompt_embeds_list = [get_prompt_embeds(chat, tokenizer, embedding_layer) for chat in chats]
outputs = llm.generate([{"prompt_embeds": embeds} for embeds in prompt_embeds_list])
print("\n[Batch Inference Outputs]")
print("-" * 30)
for i, o in enumerate(outputs):
print(f"Q{i+1}: {chats[i][0]['content']}")
print(f"A{i+1}: {o.outputs[0].text}\n")
print(f"Q{i + 1}: {chats[i][0]['content']}")
print(f"A{i + 1}: {o.outputs[0].text}\n")
print("-" * 30)

58
examples/quantization/llm-compressor/w4a8_dynamic_moe.py Normal file
View File

@@ -0,0 +1,58 @@
from llmcompressor import oneshot
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "Qwen/Qwen3-30B-A3B-Instruct-2507"
# Load model.
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
recipe = """
quant_stage:
quant_modifiers:
QuantizationModifier:
ignore: ["lm_head", "re:.*mlp.gate$"]
config_groups:
group_0:
weights:
num_bits: 8
type: int
strategy: channel
dynamic: false
symmetric: true
input_activations:
num_bits: 8
type: int
strategy: token
dynamic: true
symmetric: true
targets: ["re:.*self_attn.k_proj.*", "re:.*self_attn.o_proj.*",
"re:.*self_attn.q_proj.*", "re:.*self_attn.v_proj.*"]
group_1:
weights:
num_bits: 4
type: int
strategy: group
group_size: 128
dynamic: false
symmetric: true
input_activations:
num_bits: 8
type: int
strategy: token
dynamic: true
symmetric: true
targets: ["re:.*down_proj.*", "re:.*gate_proj.*", "re:.*up_proj.*"]
"""
# Apply quantization.
oneshot(
model=model,
recipe=recipe,
trust_remote_code_model=True,
)
# Save to disk in compressed-tensors format.
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-W4A8"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)

150
examples/quantization/llm-compressor/w8a8_int8.py Normal file
View File

@@ -0,0 +1,150 @@
import os
import torch
from compressed_tensors.quantization import QuantizationArgs, QuantizationScheme, QuantizationStrategy, QuantizationType
from datasets import load_dataset
from llmcompressor import oneshot
from llmcompressor.modifiers.awq import AWQModifier
from llmcompressor.modifiers.quantization import GPTQModifier, QuantizationModifier
from transformers import (
AutoConfig,
AutoModelForCausalLM,
AutoTokenizer,
)
W8A8_W_cha_A_ten_static_symmetric = {
"group_0": QuantizationScheme(
targets=["Linear"],
weights=QuantizationArgs(
num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.CHANNEL, symmetric=True, dynamic=False
),
input_activations=QuantizationArgs(
num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.TENSOR, symmetric=True, dynamic=False
),
),
}
# supported modifiers
MODIFIER_DICT = {
"PTQ": QuantizationModifier,
"AWQ": AWQModifier,
"GPTQ": GPTQModifier,
}
# supported schemes
SCHEMES_DICT = {
"W8A8_W_cha_A_ten_static_symmetric": W8A8_W_cha_A_ten_static_symmetric,
}
MODEL_DICT = {
"qwen3": AutoModelForCausalLM,
}
TOKENIZER_DICT = {
"qwen3": AutoTokenizer,
}
def load_environment_variables():
env_vars = {
"model_path": "Qwen/Qwen3-32B",
"export_path": "/llm-compressor/export/GPTQ/W8A8_W_cha_A_ten_static_symmetric",
"modifier": "GPTQ",
"schemes": "W8A8_W_cha_A_ten_static_symmetric",
"calib_prompt_path": "HuggingFaceH4/ultrachat_200k",
}
# verify export model path
if env_vars["export_path"] is None:
env_vars["export_path"] = env_vars["model_path"].rstrip("/") + "-" + env_vars["modifier"]
if env_vars["schemes"] is not None:
env_vars["export_path"] += "-" + env_vars["schemes"]
os.makedirs(env_vars["export_path"], exist_ok=True)
return env_vars
def load_calibration_text_dataset(calib_prompt_path, tokenizer):
# Load dataset
for f in os.listdir(calib_prompt_path):
print(f)
if any(f.lower().endswith(".jsonl") for f in os.listdir(calib_prompt_path)):
ds = load_dataset("json", data_dir=calib_prompt_path, split="validation")
elif any(f.lower().endswith(".parquet") for f in os.listdir(calib_prompt_path)):
ds = load_dataset("parquet", data_dir=calib_prompt_path, split="train[:512]")
else:
raise ValueError("Unsupported calibration file format: {}".format(calib_prompt_path.split(".")[-1]))
# Preprocess dataset
def preprocess(example):
if tokenizer.chat_template is not None:
return {"text": tokenizer.apply_chat_template(example["messages"], tokenize=False)}
else:
return {"text": example["messages"]}
# Tokenize inputs
def tokenize(sample):
return tokenizer(
sample["text"],
add_special_tokens=False,
)
ds = ds.map(preprocess)
ds = ds.map(tokenize, remove_columns=ds.column_names)
return ds
# Define a oneshot data collator for multimodal inputs.
def data_collator(batch):
assert len(batch) == 1
return {
key: torch.tensor(value, dtype=torch.bfloat16 if key == "pixel_values" else torch.long)
for key, value in batch[0].items()
}
def quantize_model(model, env_vars, dataset_dict=None):
# since the MoE gate layers are sensitive to quantization, we add them to the ignore
# list so they remain at full precision
ignore = ["lm_head", "re:.*mlp.down_proj"]
# define a llmcompressor recipe
recipe = [
MODIFIER_DICT[env_vars["modifier"]](
config_groups=SCHEMES_DICT[env_vars["schemes"]],
ignore=ignore,
),
]
# quantize the model
oneshot(
model=model,
dataset=dataset_dict,
recipe=recipe,
trust_remote_code_model=True,
)
def save_quantized_model(model, tokenizer, save_path, save_compressed=False):
model.save_pretrained(save_path, save_compressed=save_compressed)
tokenizer.save_pretrained(save_path)
if __name__ == "__main__":
# get environment variables
env_vars = load_environment_variables()
# support model type list
config = AutoConfig.from_pretrained(env_vars["model_path"], trust_remote_code=True)
model_type = config.model_type
model = MODEL_DICT[model_type].from_pretrained(env_vars["model_path"], torch_dtype="auto", trust_remote_code=True)
tokenizer = TOKENIZER_DICT[model_type].from_pretrained(env_vars["model_path"], trust_remote_code=True)
ds = load_calibration_text_dataset(env_vars["calib_prompt_path"], tokenizer)
# Quantize the model
quantize_model(model, env_vars, ds)
# save the quantized model
save_quantized_model(model, tokenizer, env_vars["export_path"], True)

82
examples/quantization/llm-compressor/w8a8_int8_dynamic.py Normal file
View File

@@ -0,0 +1,82 @@
from datasets import load_dataset
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import GPTQModifier
from llmcompressor.modifiers.smoothquant import SmoothQuantModifier
from llmcompressor.utils import dispatch_for_generation
from transformers import AutoModelForCausalLM, AutoTokenizer
# Select model and load it.
MODEL_ID = "meta-llama/Meta-Llama-3-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
# Select calibration dataset.
DATASET_ID = "HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT = "train_sft"
# Select number of samples. 512 samples is a good place to start.
# Increasing the number of samples can improve accuracy.
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQUENCE_LENGTH = 2048
# Load dataset and preprocess.
ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
ds = ds.shuffle(seed=42)
def preprocess(example):
return {
"text": tokenizer.apply_chat_template(
example["messages"],
tokenize=False,
)
}
ds = ds.map(preprocess)
# Tokenize inputs.
def tokenize(sample):
return tokenizer(
sample["text"],
padding=False,
max_length=MAX_SEQUENCE_LENGTH,
truncation=True,
add_special_tokens=False,
)
ds = ds.map(tokenize, remove_columns=ds.column_names)
# Configure algorithms. In this case, we:
# * apply SmoothQuant to make the activations easier to quantize
# * quantize the weights to int8 with GPTQ (static per channel)
# * quantize the activations to int8 (dynamic per token)
recipe = [
SmoothQuantModifier(smoothing_strength=0.8),
GPTQModifier(targets="Linear", scheme="W8A8", ignore=["lm_head"]),
]
# Apply algorithms and save to output_dir
oneshot(
model=model,
dataset=ds,
recipe=recipe,
max_seq_length=MAX_SEQUENCE_LENGTH,
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)
# Confirm generations of the quantized model look sane.
print("\n\n")
print("========== SAMPLE GENERATION ==============")
dispatch_for_generation(model)
input_ids = tokenizer("Hello my name is", return_tensors="pt").input_ids.to("npu")
output = model.generate(input_ids, max_new_tokens=100)
print(tokenizer.decode(output[0]))
print("==========================================\n\n")
# Save to disk compressed.
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-W8A8-Dynamic-Per-Token"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)

26
examples/quantization/llm-compressor/w8a8_int8_dynamic_moe.py Normal file
View File

@@ -0,0 +1,26 @@
import torch
from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from transformers import AutoModelForCausalLM, AutoTokenizer
MODEL_ID = "Qwen/Qwen3-30B-A3B-Instruct-2507"
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, dtype=torch.bfloat16, trust_remote_code=True)
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
recipe = QuantizationModifier(
targets="Linear",
scheme="INT8",
ignore=["lm_head", "re:.*mlp.gate$"],
)
oneshot(
model=model,
recipe=recipe,
trust_remote_code_model=True,
)
# Save to disk in compressed-tensors format.
SAVE_DIR = MODEL_ID.rstrip("/").split("/")[-1] + "-INT8_W8A8"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)

509
examples/rfork/rfork_planner.py Normal file
View File

@@ -0,0 +1,509 @@
#
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
# Copyright 2023 The vLLM team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# This file is a part of the vllm-ascend project.
#
"""
Standalone rfork planner mock server used by vLLM rfork seed protocol tests.
Usage:
python examples/rfork/rfork_planner.py --host 0.0.0.0 --port 1223
"""
from __future__ import annotations
import argparse
import os
import threading
import time
import uuid
from collections import defaultdict
from collections.abc import Callable, Iterable, Mapping
from contextlib import asynccontextmanager
from dataclasses import dataclass
from fastapi import APIRouter, FastAPI, Request, Response, status
@dataclass(frozen=True)
class Settings:
host: str = "0.0.0.0"
port: int = 1223
heartbeat_ttl_sec: int = 60
heartbeat_sweep_sec: int = 5
default_resource_points: int = 1
alloc_policy: str = "fifo"
def __post_init__(self) -> None:
if self.port <= 0:
raise ValueError("port must be > 0")
if self.heartbeat_ttl_sec <= 0:
raise ValueError("heartbeat_ttl_sec must be > 0")
if self.heartbeat_sweep_sec <= 0:
raise ValueError("heartbeat_sweep_sec must be > 0")
if self.default_resource_points <= 0:
raise ValueError("default_resource_points must be > 0")
if self.alloc_policy not in {"fifo", "lru"}:
raise ValueError("alloc_policy must be one of: fifo, lru")
@staticmethod
def from_env() -> Settings:
return Settings(
host=os.getenv("RFORK_MOCK_HOST", "0.0.0.0"),
port=int(os.getenv("RFORK_MOCK_PORT", "1223")),
heartbeat_ttl_sec=int(os.getenv("RFORK_MOCK_HEARTBEAT_TTL_SEC", "60")),
heartbeat_sweep_sec=int(os.getenv("RFORK_MOCK_HEARTBEAT_SWEEP_SEC", "5")),
default_resource_points=int(os.getenv("RFORK_MOCK_DEFAULT_RESOURCE_POINTS", "1")),
alloc_policy=os.getenv("RFORK_MOCK_ALLOC_POLICY", "fifo").lower(),
)
@dataclass
class SeedRecord:
seed_key: str
seed_ip: str
seed_port: int
seed_rank: int
last_heartbeat_ts: float
resource_total: int
resource_used: int = 0
@property
def identity(self) -> str:
return f"{self.seed_ip}:{self.seed_port}:{self.seed_rank}"
@property
def available_points(self) -> int:
return max(self.resource_total - self.resource_used, 0)
@dataclass
class LeaseRecord:
user_id: str
seed_key: str
seed_identity: str
allocated_points: int
leased_at: float
class Scheduler:
def __init__(self, alloc_policy: str = "fifo") -> None:
if alloc_policy not in {"fifo", "lru"}:
raise ValueError(f"unsupported alloc policy: {alloc_policy}")
self.alloc_policy = alloc_policy
def choose_seed(self, seeds: Iterable[SeedRecord]) -> SeedRecord | None:
candidates = [seed for seed in seeds if seed.available_points > 0]
if not candidates:
return None
if self.alloc_policy == "fifo":
return min(candidates, key=lambda s: (s.last_heartbeat_ts, s.identity))
return max(candidates, key=lambda s: (s.last_heartbeat_ts, s.identity))
class Store:
def __init__(
self,
*,
heartbeat_ttl_sec: int,
default_resource_points: int,
scheduler: Scheduler,
time_fn: Callable[[], float] | None = None,
) -> None:
self._lock = threading.RLock()
self._seeds: dict[str, SeedRecord] = {}
self._seeds_by_key: dict[str, set[str]] = defaultdict(set)
self._leases: dict[str, LeaseRecord] = {}
self._heartbeat_ttl_sec = heartbeat_ttl_sec
self._default_resource_points = default_resource_points
self._scheduler = scheduler
self._time = time_fn or time.time
@staticmethod
def _seed_identity(seed_ip: str, seed_port: int, seed_rank: int) -> str:
return f"{seed_ip}:{seed_port}:{seed_rank}"
def add_seed(
self,
*,
seed_key: str,
seed_ip: str,
seed_port: int,
seed_rank: int,
resource_total: int | None = None,
) -> SeedRecord:
identity = self._seed_identity(seed_ip, seed_port, seed_rank)
now = self._time()
total = self._default_resource_points if resource_total is None else max(resource_total, 1)
with self._lock:
current = self._seeds.get(identity)
if current is None:
current = SeedRecord(
seed_key=seed_key,
seed_ip=seed_ip,
seed_port=seed_port,
seed_rank=seed_rank,
last_heartbeat_ts=now,
resource_total=total,
resource_used=0,
)
self._seeds[identity] = current
self._seeds_by_key[seed_key].add(identity)
return current
if current.seed_key != seed_key:
self._seeds_by_key[current.seed_key].discard(identity)
self._seeds_by_key[seed_key].add(identity)
current.seed_key = seed_key
current.last_heartbeat_ts = now
if resource_total is not None:
current.resource_total = max(resource_total, 1)
if current.resource_used > current.resource_total:
current.resource_used = current.resource_total
return current
def get_seed(self, *, seed_key: str) -> tuple[SeedRecord, LeaseRecord] | None:
with self._lock:
self.gc_stale_seeds_locked()
seed_identities = self._seeds_by_key.get(seed_key, set())
seeds = [self._seeds[sid] for sid in seed_identities if sid in self._seeds]
selected = self._scheduler.choose_seed(seeds)
if selected is None:
return None
selected.resource_used += 1
user_id = uuid.uuid4().hex
lease = LeaseRecord(
user_id=user_id,
seed_key=seed_key,
seed_identity=selected.identity,
allocated_points=1,
leased_at=self._time(),
)
self._leases[user_id] = lease
return selected, lease
def put_seed(self, *, seed_ip: str, seed_port: int, seed_rank: int, user_id: str) -> bool:
identity = self._seed_identity(seed_ip, seed_port, seed_rank)
with self._lock:
lease = self._leases.get(user_id)
if lease is None:
return False
if lease.seed_identity != identity:
return False
seed = self._seeds.get(identity)
if seed is not None:
seed.resource_used = max(0, seed.resource_used - lease.allocated_points)
del self._leases[user_id]
return True
def gc_stale_seeds(self) -> int:
with self._lock:
return self.gc_stale_seeds_locked()
def gc_stale_seeds_locked(self) -> int:
now = self._time()
stale_ids = [
sid for sid, seed in self._seeds.items() if (now - seed.last_heartbeat_ts) > self._heartbeat_ttl_sec
]
if not stale_ids:
return 0
stale_set = set(stale_ids)
for sid in stale_ids:
seed = self._seeds.pop(sid)
self._seeds_by_key[seed.seed_key].discard(sid)
if not self._seeds_by_key[seed.seed_key]:
del self._seeds_by_key[seed.seed_key]
lease_ids = [uid for uid, lease in self._leases.items() if lease.seed_identity in stale_set]
for uid in lease_ids:
del self._leases[uid]
return len(stale_ids)
def debug_snapshot(self) -> dict[str, object]:
with self._lock:
return {
"seed_count": len(self._seeds),
"lease_count": len(self._leases),
"seeds": {
sid: {
"seed_key": s.seed_key,
"resource_total": s.resource_total,
"resource_used": s.resource_used,
"last_heartbeat_ts": s.last_heartbeat_ts,
}
for sid, s in self._seeds.items()
},
"leases": {
uid: {
"seed_identity": lease.seed_identity,
"seed_key": lease.seed_key,
"allocated_points": lease.allocated_points,
}
for uid, lease in self._leases.items()
},
}
class HeartbeatGc:
def __init__(self, store: Store, sweep_interval_sec: int) -> None:
self._store = store
self._sweep_interval_sec = max(sweep_interval_sec, 1)
self._stop = threading.Event()
self._thread: threading.Thread | None = None
def start(self) -> None:
if self._thread is not None and self._thread.is_alive():
return
self._stop.clear()
self._thread = threading.Thread(target=self._run, name="rfork-heartbeat-gc", daemon=True)
self._thread.start()
def stop(self) -> None:
self._stop.set()
if self._thread is not None:
self._thread.join(timeout=2)
def _run(self) -> None:
while not self._stop.is_set():
self._store.gc_stale_seeds()
time.sleep(self._sweep_interval_sec)
class HeaderError(ValueError):
pass
@dataclass(frozen=True)
class AddSeedHeaders:
seed_key: str
seed_ip: str
seed_port: int
seed_rank: int
seed_refcnt: int
@dataclass(frozen=True)
class GetSeedHeaders:
seed_key: str
@dataclass(frozen=True)
class PutSeedHeaders:
seed_ip: str
seed_port: int
seed_rank: int
user_id: str
def _required(headers: Mapping[str, str], key: str) -> str:
value = headers.get(key)
if value is None or value == "":
raise HeaderError(f"missing required header: {key}")
return value
def _parse_int(value: str, key: str, *, minimum: int = 0) -> int:
try:
parsed = int(value)
except ValueError as exc:
raise HeaderError(f"invalid integer header {key}: {value}") from exc
if parsed < minimum:
raise HeaderError(f"header {key} must be >= {minimum}, got {parsed}")
return parsed
def parse_add_seed_headers(headers: Mapping[str, str]) -> AddSeedHeaders:
return AddSeedHeaders(
seed_key=_required(headers, "SEED_KEY"),
seed_ip=_required(headers, "SEED_IP"),
seed_port=_parse_int(_required(headers, "SEED_PORT"), "SEED_PORT", minimum=1),
seed_rank=_parse_int(_required(headers, "SEED_RANK"), "SEED_RANK", minimum=0),
seed_refcnt=_parse_int(_required(headers, "SEED_REFCNT"), "SEED_REFCNT", minimum=0),
)
def parse_get_seed_headers(headers: Mapping[str, str]) -> GetSeedHeaders:
return GetSeedHeaders(seed_key=_required(headers, "SEED_KEY"))
def parse_put_seed_headers(headers: Mapping[str, str]) -> PutSeedHeaders:
return PutSeedHeaders(
seed_ip=_required(headers, "SEED_IP"),
seed_port=_parse_int(_required(headers, "SEED_PORT"), "SEED_PORT", minimum=1),
seed_rank=_parse_int(_required(headers, "SEED_RANK"), "SEED_RANK", minimum=0),
user_id=_required(headers, "USER_ID"),
)
def build_router(store: Store):
router = APIRouter()
@router.post("/add_seed")
def add_seed(request: Request) -> Response:
try:
parsed = parse_add_seed_headers(request.headers)
except HeaderError as err:
return Response(content=str(err), status_code=status.HTTP_400_BAD_REQUEST)
store.add_seed(
seed_key=parsed.seed_key,
seed_ip=parsed.seed_ip,
seed_port=parsed.seed_port,
seed_rank=parsed.seed_rank,
# vLLM currently sends SEED_REFCNT=0 as heartbeat metadata.
# Capacity is controlled by planner config, not by this field.
resource_total=None,
)
return Response(status_code=status.HTTP_200_OK)
@router.get("/get_seed")
def get_seed(request: Request) -> Response:
try:
parsed = parse_get_seed_headers(request.headers)
except HeaderError as err:
return Response(content=str(err), status_code=status.HTTP_400_BAD_REQUEST)
result = store.get_seed(seed_key=parsed.seed_key)
if result is None:
return Response(content="no available seed", status_code=status.HTTP_404_NOT_FOUND)
seed, lease = result
response = Response(status_code=status.HTTP_200_OK)
response.headers["SEED_IP"] = seed.seed_ip
response.headers["SEED_PORT"] = str(seed.seed_port)
response.headers["SEED_RANK"] = str(seed.seed_rank)
response.headers["USER_ID"] = lease.user_id
return response
@router.post("/put_seed")
def put_seed(request: Request) -> Response:
try:
parsed = parse_put_seed_headers(request.headers)
except HeaderError as err:
return Response(content=str(err), status_code=status.HTTP_400_BAD_REQUEST)
released = store.put_seed(
seed_ip=parsed.seed_ip,
seed_port=parsed.seed_port,
seed_rank=parsed.seed_rank,
user_id=parsed.user_id,
)
if not released:
return Response(content="lease not found", status_code=status.HTTP_404_NOT_FOUND)
return Response(status_code=status.HTTP_200_OK)
@router.get("/healthz")
def healthz() -> dict[str, str]:
return {"status": "ok"}
@router.get("/debug/snapshot")
def debug_snapshot() -> dict[str, object]:
return store.debug_snapshot()
return router
def create_app(settings: Settings):
scheduler = Scheduler(settings.alloc_policy)
store = Store(
heartbeat_ttl_sec=settings.heartbeat_ttl_sec,
default_resource_points=settings.default_resource_points,
scheduler=scheduler,
)
gc_runner = HeartbeatGc(store, settings.heartbeat_sweep_sec)
@asynccontextmanager
async def lifespan(_: FastAPI):
gc_runner.start()
try:
yield
finally:
gc_runner.stop()
app = FastAPI(title="rfork planner mock", version="0.1.0", lifespan=lifespan)
app.include_router(build_router(store))
app.state.settings = settings
app.state.store = store
app.state.gc_runner = gc_runner
return app
def _build_arg_parser() -> argparse.ArgumentParser:
defaults = Settings.from_env()
parser = argparse.ArgumentParser(description="Standalone rfork planner server")
parser.add_argument("--host", default=defaults.host, help="bind host (default: env RFORK_MOCK_HOST or 0.0.0.0)")
parser.add_argument(
"--port",
type=int,
default=defaults.port,
help="bind port (default: env RFORK_MOCK_PORT or 1223)",
)
parser.add_argument(
"--heartbeat-ttl-sec",
type=int,
default=defaults.heartbeat_ttl_sec,
help="seed heartbeat ttl in seconds (default: env RFORK_MOCK_HEARTBEAT_TTL_SEC or 60)",
)
parser.add_argument(
"--heartbeat-sweep-sec",
type=int,
default=defaults.heartbeat_sweep_sec,
help="gc sweep interval in seconds (default: env RFORK_MOCK_HEARTBEAT_SWEEP_SEC or 5)",
)
parser.add_argument(
"--default-resource-points",
type=int,
default=defaults.default_resource_points,
help="default seed capacity points (default: env RFORK_MOCK_DEFAULT_RESOURCE_POINTS or 1)",
)
parser.add_argument(
"--alloc-policy",
choices=["fifo", "lru"],
default=defaults.alloc_policy,
help="seed allocation policy (default: env RFORK_MOCK_ALLOC_POLICY or fifo)",
)
return parser
def main() -> None:
parser = _build_arg_parser()
args = parser.parse_args()
settings = Settings(
host=args.host,
port=args.port,
heartbeat_ttl_sec=args.heartbeat_ttl_sec,
heartbeat_sweep_sec=args.heartbeat_sweep_sec,
default_resource_points=args.default_resource_points,
alloc_policy=args.alloc_policy,
)
app = create_app(settings)
try:
import uvicorn
except ModuleNotFoundError as exc:
raise SystemExit("missing dependency: uvicorn. Install it with: python -m pip install uvicorn") from exc
uvicorn.run(app, host=settings.host, port=settings.port)
if __name__ == "__main__":
main()

290
examples/rl/rlhf_http_hccl.py Normal file
View File

@@ -0,0 +1,290 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
via HTTP API, with native weight syncing APIs.
Unlike rlhf.py which creates a vLLM instance programmatically, this script
assumes you have already started a vLLM server using `vllm serve`. It uses:
- OpenAI-compatible API for inference requests
- HTTP endpoints for weight transfer control plane
- HCCL for actual weight data transfer
Prerequisites:
Start a vLLM server with weight transfer enabled:
$ VLLM_SERVER_DEV_MODE=1 vllm serve Qwen/Qwen3-0.6b \
--enforce-eager \
--weight-transfer-config '{"backend": "nccl"}' \
--load-format dummy
Then run this script:
$ python rlhf_http_hccl.py
The example performs the following steps:
* Load the training model on NPU 0.
* Generate text using the vLLM server via OpenAI-compatible API. The output
is expected to be nonsense because the server is initialized with dummy weights.
* Initialize weight transfer via HTTP endpoint.
* Broadcast the real weights from the training model to the vLLM server
using HCCL.
* Generate text again to show normal output after the weight update.
"""
import requests
import torch
from openai import OpenAI
from transformers import AutoModelForCausalLM
from vllm.utils.network_utils import get_ip, get_open_port
from vllm_ascend.distributed.weight_transfer.hccl_engine import (
HCCLTrainerSendWeightsArgs,
HCCLWeightTransferEngine,
)
BASE_URL = "http://localhost:8000"
MODEL_NAME = "Qwen/Qwen3-0.6B"
def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
"""Generate completions using the OpenAI-compatible API."""
results = []
for prompt in prompts:
response = client.completions.create(
model=model,
prompt=prompt,
max_tokens=32,
temperature=0,
)
results.append(response.choices[0].text)
return results
def init_weight_transfer_engine(
base_url: str,
master_address: str,
master_port: int,
rank_offset: int,
world_size: int,
) -> None:
"""Initialize weight transfer via HTTP endpoint."""
url = f"{base_url}/init_weight_transfer_engine"
payload = {
"init_info": dict(
master_address=master_address,
master_port=master_port,
rank_offset=rank_offset,
world_size=world_size,
)
}
response = requests.post(url, json=payload, timeout=60)
response.raise_for_status()
def update_weights(
base_url: str,
names: list[str],
dtype_names: list[str],
shapes: list[list[int]],
packed: bool = False,
packed_buffer_size_bytes: int | None = None,
) -> None:
"""Update weights via HTTP endpoint."""
url = f"{base_url}/update_weights"
payload = {
"update_info": dict(
names=names,
dtype_names=dtype_names,
shapes=shapes,
packed=packed,
)
}
if packed and packed_buffer_size_bytes is not None:
payload["update_info"]["packed_buffer_size_bytes"] = packed_buffer_size_bytes
response = requests.post(url, json=payload, timeout=300)
response.raise_for_status()
def start_weight_update(base_url: str, is_checkpoint_format: bool = True) -> None:
"""Start weight update via HTTP endpoint.
Prepares the model for layerwise reload on the vLLM server side.
Must be called before update_weights.
"""
url = f"{base_url}/start_weight_update"
payload = {"is_checkpoint_format": is_checkpoint_format}
response = requests.post(url, json=payload, timeout=60)
response.raise_for_status()
def finish_weight_update(base_url: str) -> None:
"""Finish weight update via HTTP endpoint.
Finalizes layerwise reload on the vLLM server side.
Must be called after all update_weights calls are complete.
"""
url = f"{base_url}/finish_weight_update"
response = requests.post(url, timeout=60)
response.raise_for_status()
def pause_generation(base_url: str) -> None:
"""Pause generation via HTTP endpoint."""
url = f"{base_url}/pause"
response = requests.post(url, timeout=60)
response.raise_for_status()
def resume_generation(base_url: str) -> None:
"""Resume generation via HTTP endpoint."""
url = f"{base_url}/resume"
response = requests.post(url, timeout=60)
response.raise_for_status()
def get_world_size(base_url: str) -> int:
"""Get world size from the vLLM server."""
url = f"{base_url}/get_world_size"
response = requests.get(url, timeout=10)
response.raise_for_status()
return response.json()["world_size"]
def main():
# Get the inference world size from the vLLM server
inference_world_size = get_world_size(BASE_URL)
world_size = inference_world_size + 1 # +1 for the trainer
device = f"npu:{inference_world_size}"
torch.accelerator.set_device_index(device)
# Load the training model
print(f"Loading training model: {MODEL_NAME}")
train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
train_model.to(device)
# Create OpenAI client pointing to the vLLM server
client = OpenAI(
base_url=f"{BASE_URL}/v1",
api_key="EMPTY", # vLLM doesn't require an API key by default
)
# Test prompts
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Generate text before weight update. The output is expected to be nonsense
# because the server is initialized with dummy weights.
print("-" * 50)
print("Generating text BEFORE weight update (expect nonsense):")
print("-" * 50)
outputs = generate_completions(client, MODEL_NAME, prompts)
for prompt, generated_text in zip(prompts, outputs):
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
# Set up the communication channel between the training process and the
# vLLM server. The trainer is rank 0, vLLM worker(s) start at rank_offset.
master_address = get_ip()
master_port = get_open_port()
rank_offset = 1
print(f"Initializing weight transfer: master={master_address}:{master_port}")
# Initialize weight transfer on vLLM server (this is async, server will
# wait for HCCL connection)
import threading
init_thread = threading.Thread(
target=init_weight_transfer_engine,
args=(BASE_URL, master_address, master_port, rank_offset, world_size),
)
init_thread.start()
# Initialize HCCL process group on trainer side
model_update_group = HCCLWeightTransferEngine.trainer_init(
dict(
master_address=master_address,
master_port=master_port,
world_size=world_size,
),
)
# Wait for init_weight_transfer_engine to complete
init_thread.join()
# Pause generation before weight sync
pause_generation(BASE_URL)
# Start weight update (prepares layerwise reload on the vLLM server)
start_weight_update(BASE_URL)
# Collect weight metadata for the update request.
# Also track the largest tensor to auto-size the packed buffer.
names = []
dtype_names = []
shapes = []
max_tensor_bytes = 0
for name, p in train_model.named_parameters():
names.append(name)
dtype_names.append(str(p.dtype).split(".")[-1])
shapes.append(list(p.shape))
tensor_bytes = p.numel() * p.element_size()
if tensor_bytes > max_tensor_bytes:
max_tensor_bytes = tensor_bytes
# Size the packed buffer to fit the largest tensor with 128 MB headroom,
# but keep the default 1 GB when the largest tensor is smaller than that.
packed_buffer_size_bytes = max(max_tensor_bytes + 128 * 2**20, 2**30)
print(
f"Largest tensor: {max_tensor_bytes / 2**30:.2f} GiB, packed buffer: {packed_buffer_size_bytes / 2**30:.2f} GiB"
)
# Start the update_weights call in a separate thread since it will block
# waiting for HCCL broadcasts
# packed=True enables efficient batched tensor broadcasting
update_thread = threading.Thread(
target=update_weights,
args=(BASE_URL, names, dtype_names, shapes, True, packed_buffer_size_bytes),
)
update_thread.start()
# Broadcast all weights from trainer to vLLM workers
print("Broadcasting weights via HCCL...")
trainer_args = HCCLTrainerSendWeightsArgs(
group=model_update_group,
packed=True,
packed_buffer_size_bytes=packed_buffer_size_bytes,
)
HCCLWeightTransferEngine.trainer_send_weights(
iterator=train_model.named_parameters(),
trainer_args=trainer_args,
)
# Wait for update_weights to complete
update_thread.join()
# Finish weight update (finalizes layerwise reload on the vLLM server)
finish_weight_update(BASE_URL)
# Resume generation after weight sync
resume_generation(BASE_URL)
# Generate text after weight update. The output is expected to be normal
# because the real weights are now loaded.
print("-" * 50)
print("Generating text AFTER weight update:")
print("-" * 50)
outputs_updated = generate_completions(client, MODEL_NAME, prompts)
for prompt, generated_text in zip(prompts, outputs_updated):
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
if __name__ == "__main__":
main()

198
examples/rl/rlhf_http_npu_ipc.py Normal file
View File

@@ -0,0 +1,198 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Demonstrates reinforcement learning from human feedback (RLHF) using vLLM
via HTTP API, with NPU IPC-based weight syncing APIs.
Unlike rlhf_http_hccl.py which uses HCCL and can use separate NPUs, this script
uses Ascend NPU IPC which requires the training model and vLLM server to be on
the same physical NPU. Memory must be carefully managed to fit both models.
Prerequisites:
Start a vLLM server with weight transfer enabled and reduced NPU memory
utilization to leave room for the training model:
$ VLLM_SERVER_DEV_MODE=1 VLLM_ALLOW_INSECURE_SERIALIZATION=1 \
vllm serve Qwen/Qwen3-0.6b --enforce-eager \
--weight-transfer-config '{"backend": "ipc"}' \
--load-format dummy \
--gpu-memory-utilization 0.5
Then run this script:
$ python rlhf_http_npu_ipc.py
The example performs the following steps:
* Load the training model on NPU 0 (same NPU as the vLLM server).
* Generate text using the vLLM server via OpenAI-compatible API. The output
is expected to be nonsense because the server is initialized with dummy weights.
* Initialize weight transfer via HTTP endpoint (no-op for NPU IPC).
* Pause generation and broadcast the real weights from the training model to
the vLLM server using NPU IPC handles (via HTTP). The pause/resume is
handled by ``trainer_send_weights`` — it calls ``update_weights`` internally.
* Generate text again to show normal output after the weight update.
"""
import os
import requests
import torch
from openai import OpenAI
from transformers import AutoModelForCausalLM
from vllm_ascend.distributed.weight_transfer.npu_ipc_engine import (
NPUIPCTrainerSendWeightsArgs,
NPUIPCWeightTransferEngine,
)
BASE_URL = "http://localhost:8000"
MODEL_NAME = "Qwen/Qwen3-0.6B"
# Enable insecure serialization for IPC handle serialization over HTTP
os.environ["VLLM_ALLOW_INSECURE_SERIALIZATION"] = "1"
def generate_completions(client: OpenAI, model: str, prompts: list[str]) -> list[str]:
"""Generate completions using the OpenAI-compatible API."""
results = []
for prompt in prompts:
response = client.completions.create(
model=model,
prompt=prompt,
max_tokens=32,
temperature=0,
)
results.append(response.choices[0].text)
return results
def init_weight_transfer_engine(base_url: str) -> None:
"""Initialize weight transfer via HTTP endpoint (no-op for NPU IPC)."""
url = f"{base_url}/init_weight_transfer_engine"
payload: dict[str, dict] = {"init_info": {}}
response = requests.post(url, json=payload, timeout=60)
response.raise_for_status()
def start_weight_update(base_url: str, is_checkpoint_format: bool = True) -> None:
"""Start weight update via HTTP endpoint.
Prepares the model for layerwise reload on the vLLM server side.
Must be called before update_weights.
"""
url = f"{base_url}/start_weight_update"
payload = {"is_checkpoint_format": is_checkpoint_format}
response = requests.post(url, json=payload, timeout=60)
response.raise_for_status()
def finish_weight_update(base_url: str) -> None:
"""Finish weight update via HTTP endpoint.
Finalizes layerwise reload on the vLLM server side.
Must be called after all update_weights calls are complete.
"""
url = f"{base_url}/finish_weight_update"
response = requests.post(url, timeout=60)
response.raise_for_status()
def pause_generation(base_url: str) -> None:
"""Pause generation via HTTP endpoint."""
url = f"{base_url}/pause"
response = requests.post(url, timeout=60)
response.raise_for_status()
def resume_generation(base_url: str) -> None:
"""Resume generation via HTTP endpoint."""
url = f"{base_url}/resume"
response = requests.post(url, timeout=60)
response.raise_for_status()
def main():
# NPU IPC requires the training model to be on the same NPU as the vLLM server.
# The server should be started on NPU 0 with reduced memory utilization.
device = "npu:0"
torch.accelerator.set_device_index(device)
# Load the training model on the same NPU as the server.
# Use bfloat16 to reduce memory footprint.
print(f"Loading training model: {MODEL_NAME} on {device}")
print(
"Note: Ensure the vLLM server was started with --gpu-memory-utilization 0.5 "
"or lower to leave room for the training model."
)
train_model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, dtype=torch.bfloat16)
train_model.to(device)
train_model.eval()
# Create OpenAI client pointing to the vLLM server
client = OpenAI(
base_url=f"{BASE_URL}/v1",
api_key="EMPTY",
)
# Test prompts
prompts = [
"Hello, my name is",
"The president of the United States is",
"The capital of France is",
"The future of AI is",
]
# Generate text before weight update. The output is expected to be nonsense
# because the server is initialized with dummy weights.
print("-" * 50)
print("Generating text BEFORE weight update (expect nonsense):")
print("-" * 50)
outputs = generate_completions(client, MODEL_NAME, prompts)
for prompt, generated_text in zip(prompts, outputs):
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
# Initialize weight transfer on vLLM server (no-op for NPU IPC)
print("Initializing weight transfer (NPU IPC backend)...")
init_weight_transfer_engine(BASE_URL)
# Pause generation before weight sync
pause_generation(BASE_URL)
# Start weight update (prepares layerwise reload on the vLLM server)
start_weight_update(BASE_URL)
# Send weights via NPU IPC handles using HTTP mode.
# trainer_send_weights internally collects all parameters,
# creates IPC handles, and POSTs them to /update_weights.
print("Broadcasting weights via NPU IPC (HTTP)...")
trainer_args = NPUIPCTrainerSendWeightsArgs(send_mode="http", url=BASE_URL)
NPUIPCWeightTransferEngine.trainer_send_weights(
iterator=train_model.named_parameters(),
trainer_args=trainer_args,
)
# Finish weight update (finalizes layerwise reload on the vLLM server)
finish_weight_update(BASE_URL)
# Resume generation after weight sync
resume_generation(BASE_URL)
# Generate text after weight update. The output is expected to be normal
# because the real weights are now loaded.
print("-" * 50)
print("Generating text AFTER weight update:")
print("-" * 50)
outputs_updated = generate_completions(client, MODEL_NAME, prompts)
for prompt, generated_text in zip(prompts, outputs_updated):
print(f"Prompt: {prompt!r}\nGenerated text: {generated_text!r}")
print("-" * 50)
# Note: The training model and IPC handles remain in memory.
# In a real RLHF training loop, you would update the training model
# and create new IPC handles for each weight update.
if __name__ == "__main__":
main()

8
examples/run_dp_server.sh
View File

@@ -5,10 +5,9 @@ export TP_SOCKET_IFNAME="eth0"
export HCCL_SOCKET_IFNAME="eth0"
export OMP_PROC_BIND=false
export OMP_NUM_THREADS=100
export OMP_NUM_THREADS=10
export VLLM_USE_V1=1
export VLLM_USE_MODELSCOPE=true
export VLLM_USE_MODELSCOPE=True
export ASCEND_LAUNCH_BLOCKING=0
@@ -28,5 +27,4 @@ vllm serve Qwen/Qwen1.5-MoE-A2.7B \
--max-num-batched-tokens 4096 \
--gpu-memory-utilization 0.9 \
--trust-remote-code \
--enforce-eager \
--additional-config '{"ascend_scheduler_config":{"enabled":true},"torchair_graph_config":{"enabled":false, "use_cached_graph":false}}'
--enforce-eager

287
examples/save_sharded_state_310.py Normal file
View File

@@ -0,0 +1,287 @@
#
# Copyright (c) 2026 Huawei Technologies Co., Ltd. All Rights Reserved.
# This file is a part of the vllm-ascend project.
# Adapted from vllm-project/vllm/examples/offline_inference/save_sharded_state.py
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
"""
Saves each worker's model state dict directly to a checkpoint, which enables a
fast load path for large tensor-parallel models where each worker only needs to
read its own shard rather than the entire checkpoint.
Sparse-Compress-Quantization state dict could also be saved via this script.
Example usage:
python save_sharded_state_310.py \
--model /path/to/load \
--tensor-parallel-size 8 \
--output /path/to/save \
--enable-compress \
--compress-process-num 8 \
--enforce-eager \
--dtype float16 \
--quantization ascend
Then, the model can be loaded with
llm = LLM(
model="/path/to/save",
load_format="sharded_state",
tensor_parallel_size=8,
quantization="ascend",
)
"""
import dataclasses
import json
import multiprocessing as mp
import os
import shutil
from pathlib import Path
import torch
from vllm import LLM, EngineArgs
from vllm.distributed.parallel_state import destroy_distributed_environment, destroy_model_parallel
from vllm.utils.argparse_utils import FlexibleArgumentParser
SUPPORTED_COMPRESS_QUANT_TYPE = ["W8A8S", "W16A16S"]
DEFAULT_PATTERN = "model-rank-{rank}-part-{part}.safetensors"
QUANTIZATION_UPDATE_MAP = {"W8A8S": "W8A8SC", "W16A16S": "W16A16SC"}
class FileHandler:
@staticmethod
def validate_path(path: str, must_exist: bool = True, check_writable: bool = False) -> Path:
"""
Comprehensive path validation.
- Checks existence
- Checks write permissions for the target or its parent
"""
p = Path(path)
if must_exist and not p.exists():
raise FileNotFoundError(f"Error: Path '{path}' does not exist.")
if check_writable:
# Check the directory itself if it exists, otherwise check the parent
target = p if p.exists() else p.parent
if not os.access(target, os.W_OK):
raise PermissionError(f"Permission Denied: No write access to '{target}'.")
return p
@staticmethod
def safe_copy(src: Path, dst: Path):
"""Copies files or directories with permission handling."""
try:
if src.is_dir():
# dirs_exist_ok=True prevents errors if the destination directory exists
shutil.copytree(src, dst, dirs_exist_ok=True)
else:
# copy2 preserves metadata (timestamps, permissions)
shutil.copy2(src, dst)
except (PermissionError, OSError) as e:
print(f"Warning: Failed to copy {src} due to: {e}")
def clean_up():
"""Clean up VLLM resources"""
destroy_model_parallel()
destroy_distributed_environment()
torch.npu.empty_cache()
def parse_args():
parser = FlexibleArgumentParser()
EngineArgs.add_cli_args(parser)
parser.add_argument("--output", "-o", required=True, type=str, help="path to output checkpoint")
parser.add_argument(
"--enable-compress",
action="store_true",
)
parser.add_argument(
"--compress-process-num",
type=int,
default=1,
)
return parser.parse_args()
def get_quant_description(json_file: str) -> dict:
"""
Extract quantization description from JSON configuration file.
Args:
json_file: Path to the JSON configuration file
Returns:
dict: Quantization descriptor dictionary
Raises:
FileNotFoundError: If the JSON file does not exist
RuntimeError: If JSON parsing fails or required keys are missing
"""
config_path = Path(json_file)
if not config_path.exists():
raise FileNotFoundError(f"Model configuration file not found: {json_file}")
try:
with config_path.open("r", encoding="utf-8") as file:
quant_desc = json.load(file)
except json.JSONDecodeError as e:
raise RuntimeError(f"Invalid JSON format in {json_file}: {e}")
return quant_desc
def update_quant_description(ori_json_file: str, target_json_file: str) -> None:
"""
Update quantization types in JSON configuration file based on update mapping.
Args:
ori_json_file: Path to the JSON configuration file
target_json_file: Path to the JSON configuration file to be saved
Raises:
FileNotFoundError: If the JSON file does not exist
RuntimeError: If JSON parsing fails or required keys are missing
"""
config_path = Path(ori_json_file)
try:
with config_path.open("r", encoding="utf-8") as file:
json_data = json.load(file)
except (FileNotFoundError, json.JSONDecodeError) as e:
raise RuntimeError(f"Failed to read configuration file {ori_json_file}: {e}")
original_quant_type = json_data.get("model_quant_type")
if not original_quant_type or original_quant_type not in QUANTIZATION_UPDATE_MAP:
raise RuntimeError(
f"Cannot update quantization type. "
f"Original type '{original_quant_type}' not found or not supported for update in {ori_json_file}."
)
updated_quant_type = QUANTIZATION_UPDATE_MAP[original_quant_type]
updated_config = {"model_quant_type": updated_quant_type, "version": "1.0.0"}
for key, value in json_data.items():
if key.endswith(".weight") and value == original_quant_type:
updated_config[key] = updated_quant_type
elif key not in ("model_quant_type", "version"):
updated_config[key] = value
try:
new_file_path = Path(target_json_file)
with new_file_path.open("w", encoding="utf-8") as file:
json.dump(updated_config, file, indent=2, ensure_ascii=False)
os.remove(ori_json_file)
except OSError as e:
raise RuntimeError(f"Failed to write updated configuration to {target_json_file}: {e}")
def weight_compress_worker(file_path: str, quant_desc: dict, process_num: int) -> bool:
"""
Worker logic for multiprocessing.
Note: Imports are inside the worker to save memory in the main process.
Returns:
bool: True if processing succeeded, False otherwise.
"""
import safetensors
import safetensors.torch
from msmodelslim.pytorch.weight_compression import CompressConfig, Compressor
p = Path(file_path)
if not p.exists():
print(f"Error: File not found, failed to compress: {file_path}")
return False
try:
state_dict = safetensors.torch.load_file(str(p))
compress_config = CompressConfig(
do_pseudo_sparse=False,
sparse_ratio=1,
is_debug=True,
record_detail_root=str(p.parent),
multiprocess_num=process_num,
)
compressor = Compressor(compress_config, weight=state_dict, quant_model_description=quant_desc)
compressor.run()
if p.exists():
os.remove(p)
compressor.export_safetensors(str(p.parent), safetensors_name=p.name)
return True
except Exception as e:
print(f"Error processing Rank file {file_path}: {e}")
return False
def main(args):
# 1. Initial Validation
# Validate early so the script doesn't fail after hours of inference
output_dir = FileHandler.validate_path(args.output, must_exist=False, check_writable=True)
model_dir = FileHandler.validate_path(args.model, must_exist=True)
# 2. Run VLLM Engine and save sharded states
engine_args = EngineArgs.from_cli_args(args)
llm = LLM(**dataclasses.asdict(engine_args))
output_dir.mkdir(parents=True, exist_ok=True)
llm.llm_engine.engine_core.save_sharded_state(path=str(output_dir))
del llm
clean_up()
# 3. Migrate Metadata (Excluding large weights)
for item in model_dir.iterdir():
if item.suffix not in (".bin", ".pt", ".safetensors"):
FileHandler.safe_copy(item, output_dir / item.name)
# 4. Compression Logic
parameters_map_fpath = output_dir / "parameters_type_map.json"
if args.enable_compress:
quant_desc_file = output_dir / "quant_model_description.json"
backup_quant_desc_file = output_dir / "ori_quant_model_description.json"
if quant_desc_file.exists():
os.rename(str(quant_desc_file), str(backup_quant_desc_file))
quant_desc = get_quant_description(str(parameters_map_fpath))
quant_type = quant_desc["model_quant_type"]
if quant_type in SUPPORTED_COMPRESS_QUANT_TYPE:
# TODO: Implement w16a16sc
if quant_type == "W16A16S":
raise NotImplementedError("W16A16SC is not supported yet.")
tasks = []
for i in range(args.tensor_parallel_size):
file_name = DEFAULT_PATTERN.format(rank=i, part="0")
full_path = output_dir / file_name
p = mp.Process(
target=weight_compress_worker, args=(str(full_path), quant_desc, args.compress_process_num)
)
tasks.append(p)
p.start()
for p in tasks:
p.join()
update_quant_description(str(backup_quant_desc_file), str(quant_desc_file))
print("Compression completed successfully.")
else:
print(f"Skipping compression: Unsupported type {quant_type}")
if parameters_map_fpath.exists():
os.remove(parameters_map_fpath)
if __name__ == "__main__":
args = parse_args()
main(args)