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

style: format Rust workspace

Browse Source
This commit is contained in:
Gahow Wang
2026-06-18 18:11:58 +08:00
parent 013465fc06
commit 531cd3fe08
57 changed files with 4045 additions and 1204 deletions
Download Patch File Download Diff File Expand all files Collapse all files

100
crates/xserv-server/src/pp_engine.rs
View File

@@ -15,15 +15,15 @@
use std::ffi::c_void;
use std::path::{Path, PathBuf};
use std::sync::mpsc;
use std::sync::Arc;
use std::sync::mpsc;
use std::thread;
use half::bf16;
use xserv_distributed::{PpContext, UniqueId};
use xserv_model::loader;
use xserv_model::sampling::SamplingParams;
use xserv_model::{sample, ModelConfig, PagedKVCache, Qwen3, BLOCK_SIZE};
use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, sample};
use xserv_tensor::{DType, Device, Tensor};
use xserv_tokenizer::Tokenizer;
@@ -38,9 +38,16 @@ enum PpCommand {
Free(usize),
/// Receive `[n_tokens, hidden]` from the previous stage, run this stage's
/// layers; if last stage, sample with `sampling` and return the token.
Prefill { n_tokens: usize, slot: usize, sampling: SamplingParams },
Prefill {
n_tokens: usize,
slot: usize,
sampling: SamplingParams,
},
/// Receive `[1, hidden]`, run this stage's layers; last stage samples.
Decode { slot: usize, sampling: SamplingParams },
Decode {
slot: usize,
sampling: SamplingParams,
},
Shutdown,
}
@@ -76,9 +83,21 @@ fn build_stage(
let max_blocks_per_seq = max_seq_len.div_ceil(BLOCK_SIZE);
let total_blocks = max_blocks_per_seq + 8; // v1 serial: one active sequence
let cache = PagedKVCache::new(
&stage_config, total_blocks, 0, 4, max_blocks_per_seq, DType::BF16, device,
&stage_config,
total_blocks,
0,
4,
max_blocks_per_seq,
DType::BF16,
device,
);
StageCtx { model, cache, pp, hidden: config.hidden(), device }
StageCtx {
model,
cache,
pp,
hidden: config.hidden(),
device,
}
}
/// Allocate a zeroed `[n, hidden]` device tensor and receive into it from `peer`.
@@ -110,7 +129,15 @@ fn worker_loop(
ack_tx: mpsc::Sender<()>,
token_tx: mpsc::Sender<u32>,
) {
let mut sc = build_stage(&model_dir, &config, stage, world, stage as u32, max_seq_len, id);
let mut sc = build_stage(
&model_dir,
&config,
stage,
world,
stage as u32,
max_seq_len,
id,
);
let is_last = stage == world - 1;
let prev = stage - 1;
let next = stage + 1;
@@ -125,7 +152,11 @@ fn worker_loop(
sc.cache.free_sequence(slot);
let _ = ack_tx.send(());
}
PpCommand::Prefill { n_tokens, slot, sampling } => {
PpCommand::Prefill {
n_tokens,
slot,
sampling,
} => {
let x = recv_hidden(&sc, n_tokens, prev);
let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
if is_last {
@@ -155,7 +186,12 @@ fn worker_loop(
/// Run the PP coordinator (stage 0) on the calling thread. Spawns worker stages
/// 1..world and consumes generation requests from `rx`.
pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Receiver<GenerateRequest>) {
pub fn run_pp(
model_dir: &Path,
world: usize,
max_seq_len: usize,
rx: mpsc::Receiver<GenerateRequest>,
) {
assert!(world >= 2, "run_pp requires world >= 2");
let config = ModelConfig::from_file(&model_dir.join("config.json"));
assert!(
@@ -179,7 +215,17 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
let model_dir = model_dir.to_path_buf();
let config = config.clone();
thread::spawn(move || {
worker_loop(stage, world, id, model_dir, config, max_seq_len, ctx_rx, ack_tx, token_tx);
worker_loop(
stage,
world,
id,
model_dir,
config,
max_seq_len,
ctx_rx,
ack_tx,
token_tx,
);
});
}
@@ -207,11 +253,14 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
wait_acks(&ack_rx);
// Prefill: embed prompt, run stage-0 layers, push hidden into the pipe.
broadcast(&cmd_txs, PpCommand::Prefill {
n_tokens: req.prompt_tokens.len(),
slot,
sampling: req.sampling.clone(),
});
broadcast(
&cmd_txs,
PpCommand::Prefill {
n_tokens: req.prompt_tokens.len(),
slot,
sampling: req.sampling.clone(),
},
);
let x = sc.model.embed(&req.prompt_tokens);
let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
send_hidden(&sc, &x, next_peer);
@@ -228,7 +277,13 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
if generated >= req.max_tokens {
break "length";
}
broadcast(&cmd_txs, PpCommand::Decode { slot, sampling: req.sampling.clone() });
broadcast(
&cmd_txs,
PpCommand::Decode {
slot,
sampling: req.sampling.clone(),
},
);
let x = sc.model.embed(&[next]);
let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache);
send_hidden(&sc, &x, next_peer);
@@ -239,9 +294,14 @@ pub fn run_pp(model_dir: &Path, world: usize, max_seq_len: usize, rx: mpsc::Rece
let tail = tokenizer.flush_decode_stream(&mut decode_buf);
if !tail.is_empty() {
let _ = req.sender.blocking_send(GenerateEvent::Token { id: next, text: tail });
let _ = req.sender.blocking_send(GenerateEvent::Token {
id: next,
text: tail,
});
}
let _ = req.sender.blocking_send(GenerateEvent::Done { finish_reason: finish.to_string() });
let _ = req.sender.blocking_send(GenerateEvent::Done {
finish_reason: finish.to_string(),
});
broadcast(&cmd_txs, PpCommand::Free(slot));
sc.cache.free_sequence(slot);
@@ -258,6 +318,8 @@ fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &
}
let text = tokenizer.decode_token_stream(token_id, buf);
if !text.is_empty() {
let _ = req.sender.blocking_send(GenerateEvent::Token { id: token_id, text });
let _ = req
.sender
.blocking_send(GenerateEvent::Token { id: token_id, text });
}
}