gahow/xserv
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

server: pipeline-parallel HTTP engine (--pp N)

Browse Source 
pp_engine::run_pp: stage-0 coordinator (scheduler/tokenizer/sampling +
stop logic) on the calling thread, worker stage threads for 1..P. Each
step the coordinator embeds + runs its layers, then the hidden state is
handed stage->stage over NCCL P2P; the last stage samples and returns
the token to stage 0 over an in-process channel. v1 is serial (one
request, one token/step) — correctness first; throughput via microbatch
overlap is future work.

main: wire --pp N (mutually exclusive with --tp).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
Gahow Wang
2026-05-29 18:45:52 +08:00
parent da3aaa134a
commit 824cc58daa
2 changed files with 280 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

18
crates/xserv-server/src/main.rs
View File

@@ -1,5 +1,6 @@
mod api;
mod engine;
mod pp_engine;
mod tp_engine;
use axum::{routing::{get, post}, Extension, Router};
@@ -19,7 +20,7 @@ pub struct AppState {
async fn main() {
let args: Vec<String> = std::env::args().collect();
if args.len() < 2 {
eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N]");
eprintln!("Usage: xserv-server <model-dir> [--port PORT] [--max-batch N] [--max-seq-len N] [--swap-space-gb N] [--tp N] [--pp N]");
std::process::exit(1);
}
@@ -52,6 +53,16 @@ async fn main() {
.and_then(|s| s.parse().ok())
.unwrap_or(1)
.max(1);
let pp: usize = args.iter()
.position(|a| a == "--pp")
.and_then(|i| args.get(i + 1))
.and_then(|s| s.parse().ok())
.unwrap_or(1)
.max(1);
if tp > 1 && pp > 1 {
eprintln!("--tp and --pp cannot be combined yet (2D TP×PP is future work)");
std::process::exit(1);
}
let model_config = ModelConfig::from_file(&model_dir.join("config.json"));
let model_max_seq_len = model_config.max_seq_len();
if model_max_seq_len == 0 {
@@ -76,7 +87,10 @@ async fn main() {
let model_dir_clone = model_dir.clone();
std::thread::spawn(move || {
if tp <= 1 {
if pp > 1 {
// Pipeline-parallel path: stage-0 coordinator + worker stage threads.
pp_engine::run_pp(&model_dir_clone, pp, max_seq_len, rx);
} else if tp <= 1 {
let mut engine = engine::Engine::load_with_swap(&model_dir_clone, max_batch, max_seq_len, swap_space_gb);
engine.run(rx);
} else {