mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-30 15:57:20 +08:00
8e488bbd55
A prompt-only engine tray (PE) and a decode-only engine tray (CE) behind the OpenAI-compatible frontend from serve/openai_server.py: - serve/pd_server.py: --role pe / --role ce. The CE frontend renders the prompt, picks a CE rank (affinity or least-loaded routing) and POSTs /prefill to the PE. The PE prefills with the host prefix cache on, exports the prompt's KV blocks and the full SWA state in the pe_do_save_kvcache hook, streams the first sampled token back in the HTTP response and pushes the KV bytes to the CE rank over TCP. The CE rank admits the sequence into decode via Scheduler.add_prompted_sequence and reclaims PE sequences through Scheduler.free_finished_seqs. - deepinfer/hicache/pd_transfer.py: export/admit of KV pages and state through the hicache page-copy kernels. The CE keeps the pinned receive buffers alive until the admit copy kernels finish (the kernels read them by raw pointer; the caching host allocator otherwise reused the memory while the copy was queued -> cudaErrorLaunchFailure). - serve/pd_run.sh start/stop/logs over ssh; serve/check_pd.py compares outputs against a single-engine server token by token; serve/nccl_xtray_test.py reports the cross-host NCCL transport. Measured on 4 GB300 (PE) + 4 GB300 (CE) against one 4-GPU tray with the AgentX replay, outputs identical: ITL p90 stays within 8.6-15 ms from conc 8 to 256 (single tray 12-98 ms), P90 interactivity 66-117 tok/s/user where the single tray falls to 10; per total chip PD is behind up to conc 64 and ahead at 256 (doc/benchmark_gb300.md). Depends on the host prefix cache (#144), the OpenAI server (#147) and the scheduler bindings (add_prompted_sequence / free_finished_seqs), which are included as the first three commits. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01LFPG29daR1FhchJBTjiEKr
50 lines
2.7 KiB
Bash
Executable File
50 lines
2.7 KiB
Bash
Executable File
#!/bin/bash
|
|
# Start/stop the prefill (PE) and decode (CE) trays of serve/pd_server.py over ssh.
|
|
#
|
|
# serve/pd_run.sh start <pe-host> <ce-host> [extra pd_server args...]
|
|
# serve/pd_run.sh stop <pe-host> <ce-host>
|
|
# serve/pd_run.sh log <host> # tail the tray's log
|
|
#
|
|
# Ports: PE HTTP 8889, CE HTTP (user-facing) 8888, CE KV receivers 8900+rank. Logs: ~/pd_<role>.log on each tray.
|
|
set -euo pipefail
|
|
CMD="${1:?start|stop|log}"; shift
|
|
MODEL="${DEEPINFER_HF_MODEL_DIR:-$HOME/models/DeepSeek-V4.1-Flash-1017}"
|
|
SERVED=deepseek-ai/DeepSeek-V4.1-Flash-1017
|
|
ENV='cd ~/DeepInfer && source env.sh && export CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES:-0,1,2,3}'
|
|
|
|
launch() { # host role args...
|
|
local host=$1 role=$2; shift 2
|
|
# ssh -n + nohup with every fd redirected, otherwise the remote session stays attached to the server.
|
|
ssh -n "$host" "$ENV && rm -f ~/pd_$role.pgid && nohup setsid python -u serve/pd_server.py --role $role --model $MODEL --served-model-name $SERVED $* > ~/pd_$role.log 2>&1 < /dev/null & echo \$! > ~/pd_$role.pgid; echo started $role on \$(hostname) pgid \$(cat ~/pd_$role.pgid)"
|
|
}
|
|
wait_ready() { # host port
|
|
local i
|
|
for i in $(seq 200); do
|
|
if curl -fs "http://$1:$2/health" >/dev/null 2>&1; then echo "$1:$2 ready"; return 0; fi
|
|
sleep 3
|
|
done
|
|
echo "$1:$2 did not become ready" >&2; return 1
|
|
}
|
|
|
|
case "$CMD" in
|
|
start)
|
|
PE=${1:?pe-host}; CE=${2:?ce-host}; shift 2
|
|
launch "$PE" pe --port 8889 "$@"
|
|
launch "$CE" ce --port 8888 --kv-port 8900 --kv-host "$CE" --pe-url "http://$PE:8889" "$@"
|
|
wait_ready "$PE" 8889; wait_ready "$CE" 8888
|
|
;;
|
|
stop)
|
|
for spec in "${1:?pe-host}:pe" "${2:?ce-host}:ce"; do
|
|
h=${spec%%:*}; r=${spec##*:}
|
|
# Kill by pid list, excluding the remote shell itself: a pkill -f pattern would match this very
|
|
# command line (it names pd_server.py) and kill the shell before anything else runs.
|
|
# Frontend + its spawned engine workers, SIGKILL (uvicorn's graceful shutdown otherwise keeps the
|
|
# children, and the GPUs, alive for minutes); then wait until the GPUs are actually free.
|
|
ssh -n "$h" "for p in \$(pgrep -f '[p]d_server.py --role $r|multiprocessing.spawn|multiprocessing.resource_tracker'); do [ \"\$p\" != \"\$\$\" ] && kill -9 \"\$p\" 2>/dev/null; done; rm -f ~/pd_$r.pgid; for i in \$(seq 60); do [ \"\$(nvidia-smi --query-gpu=memory.used --format=csv,noheader,nounits | sort -n | tail -1)\" -lt 1000 ] && break; sleep 3; done; echo stopped $r on \$(hostname), gpu mem \$(nvidia-smi --query-gpu=memory.used --format=csv,noheader | tr '\\n' ' ')" || true
|
|
done
|
|
;;
|
|
log)
|
|
ssh "${1:?host}" 'tail -n 40 ~/pd_*.log'
|
|
;;
|
|
esac
|