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

Add Qwen3.6 MoE inference support

Browse Source
This commit is contained in:
Gahow Wang
2026-07-13 20:24:41 +08:00
parent 588bfd9df3
commit a2de146fb6
27 changed files with 3153 additions and 149 deletions
Download Patch File Download Diff File Expand all files Collapse all files

165
crates/xserv-server/src/api.rs
View File

@@ -7,6 +7,7 @@ use serde::{Deserialize, Serialize};
use std::convert::Infallible;
use std::path::Path;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use tokio_stream::StreamExt;
use tokio_stream::wrappers::ReceiverStream;
use uuid::Uuid;
@@ -25,11 +26,35 @@ pub struct ChatRequest {
#[serde(default)]
pub stream: Option<bool>,
#[serde(default)]
pub stream_options: Option<StreamOptions>,
#[serde(default)]
pub temperature: Option<f32>,
#[serde(default)]
pub top_k: Option<usize>,
#[serde(default)]
pub top_p: Option<f32>,
#[serde(default)]
pub presence_penalty: Option<f32>,
#[serde(default)]
pub repetition_penalty: Option<f32>,
#[serde(default)]
pub chat_template_kwargs: Option<ChatTemplateKwargs>,
#[serde(default)]
pub enable_thinking: Option<bool>,
}
#[derive(Deserialize)]
pub struct StreamOptions {
#[serde(default)]
pub include_usage: bool,
}
#[derive(Deserialize, Default)]
pub struct ChatTemplateKwargs {
#[serde(default)]
pub enable_thinking: Option<bool>,
#[serde(default)]
pub preserve_thinking: Option<bool>,
}
#[derive(Deserialize, Serialize, Clone)]
@@ -102,12 +127,17 @@ impl ChatTemplate {
}
}
pub fn render(&self, messages: &[Message]) -> String {
pub fn render(
&self,
messages: &[Message],
enable_thinking: Option<bool>,
preserve_thinking: Option<bool>,
) -> String {
if self.source.is_empty() {
return build_prompt_hardcoded(messages, &self.model_type);
}
match self.render_jinja(messages) {
match self.render_jinja(messages, enable_thinking, preserve_thinking) {
Ok(prompt) => prompt,
Err(e) => {
eprintln!("[chat-template] Jinja render error: {e}, falling back to hardcoded");
@@ -116,7 +146,12 @@ impl ChatTemplate {
}
}
fn render_jinja(&self, messages: &[Message]) -> Result<String, minijinja::Error> {
fn render_jinja(
&self,
messages: &[Message],
enable_thinking: Option<bool>,
preserve_thinking: Option<bool>,
) -> Result<String, minijinja::Error> {
let mut env = minijinja::Environment::new();
// Register custom functions the template may call.
@@ -127,6 +162,42 @@ impl ChatTemplate {
env.add_filter("startswith", |s: String, prefix: String| -> bool {
s.starts_with(&prefix)
});
env.set_unknown_method_callback(|_, value, method, args| {
use minijinja::value::from_args;
use minijinja::{Error, ErrorKind, Value};
let Some(value) = value.as_str() else {
return Err(Error::from(ErrorKind::UnknownMethod));
};
match method {
"startswith" => {
let (prefix,): (String,) = from_args(args)?;
Ok(Value::from(value.starts_with(&prefix)))
}
"endswith" => {
let (suffix,): (String,) = from_args(args)?;
Ok(Value::from(value.ends_with(&suffix)))
}
"split" => {
let (separator,): (String,) = from_args(args)?;
Ok(Value::from_serialize(
value
.split(&separator)
.map(str::to_owned)
.collect::<Vec<_>>(),
))
}
"rstrip" => {
let (chars,): (String,) = from_args(args)?;
Ok(Value::from(value.trim_end_matches(|c| chars.contains(c))))
}
"lstrip" => {
let (chars,): (String,) = from_args(args)?;
Ok(Value::from(value.trim_start_matches(|c| chars.contains(c))))
}
_ => Err(Error::from(ErrorKind::UnknownMethod)),
}
});
env.add_template("chat", &self.source)?;
let tmpl = env.get_template("chat")?;
@@ -136,6 +207,8 @@ impl ChatTemplate {
add_generation_prompt => true,
bos_token => "",
eos_token => "",
enable_thinking => enable_thinking,
preserve_thinking => preserve_thinking,
};
tmpl.render(ctx)
@@ -256,8 +329,12 @@ fn build_prompt_gpt_oss(messages: &[Message]) -> String {
// HTTP handlers
// ---------------------------------------------------------------------------
pub async fn health() -> &'static str {
"ok"
pub async fn health(Extension(state): Extension<Arc<AppState>>) -> Response {
if state.engine_ready.load(Ordering::Acquire) {
(StatusCode::OK, "ok").into_response()
} else {
(StatusCode::SERVICE_UNAVAILABLE, "loading").into_response()
}
}
pub async fn list_models(Extension(state): Extension<Arc<AppState>>) -> Json<ModelsResponse> {
@@ -291,7 +368,10 @@ async fn chat_non_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
return response;
}
let prompt = state.chat_template.render(&req.messages);
let (enable_thinking, preserve_thinking) = template_options(&req);
let prompt = state
.chat_template
.render(&req.messages, enable_thinking, preserve_thinking);
let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
let prompt_token_count = prompt_tokens.len();
@@ -363,18 +443,26 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
return response;
}
let prompt = state.chat_template.render(&req.messages);
let (enable_thinking, preserve_thinking) = template_options(&req);
let prompt = state
.chat_template
.render(&req.messages, enable_thinking, preserve_thinking);
let prompt_tokens = state.engine_tokenizer.lock().unwrap().encode(&prompt);
let prompt_token_count = prompt_tokens.len();
let include_usage = req
.stream_options
.as_ref()
.is_some_and(|options| options.include_usage);
let max_seq_len = state.max_seq_len;
if prompt_tokens.len() >= max_seq_len {
if prompt_token_count >= max_seq_len {
return bad_request(format!(
"prompt is {} tokens, exceeds max_seq_len {}",
prompt_tokens.len(),
prompt_token_count,
max_seq_len
));
}
let max_tokens = req.max_tokens.min(max_seq_len - prompt_tokens.len());
let max_tokens = req.max_tokens.min(max_seq_len - prompt_token_count);
let (engine_tx, engine_rx) = tokio::sync::mpsc::channel::<GenerateEvent>(64);
let gen_req = GenerateRequest {
@@ -393,10 +481,12 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
tokio::spawn(async move {
let mut engine_stream = ReceiverStream::new(engine_rx);
let mut first = true;
let mut completion_token_count = 0usize;
while let Some(event) = engine_stream.next().await {
match event {
GenerateEvent::Token { text, .. } => {
completion_token_count += 1;
if first {
// First chunk: role announcement
let chunk =
@@ -422,6 +512,16 @@ fn chat_stream(state: Arc<AppState>, req: ChatRequest) -> Response {
let fr = normalize_finish_reason(&finish_reason);
let chunk = make_chunk(&id, &model_name, created, None, None, fr);
let _ = sse_tx.send(Ok(Event::default().data(chunk))).await;
if include_usage {
let usage = make_usage_chunk(
&id,
&model_name,
created,
prompt_token_count,
completion_token_count,
);
let _ = sse_tx.send(Ok(Event::default().data(usage))).await;
}
let _ = sse_tx
.send(Ok(Event::default().data("[DONE]".to_string())))
.await;
@@ -464,6 +564,18 @@ fn validate_request(req: &ChatRequest, model_name: &str) -> Option<Response> {
return Some(bad_request("top_k must be <= 1_000_000"));
}
}
if let Some(p) = req.presence_penalty {
if !p.is_finite() || !(-2.0..=2.0).contains(&p) {
return Some(bad_request("presence_penalty must be in [-2, 2]"));
}
}
if let Some(p) = req.repetition_penalty {
if !p.is_finite() || p <= 0.0 {
return Some(bad_request(
"repetition_penalty must be a finite value greater than 0",
));
}
}
None
}
@@ -546,6 +658,28 @@ fn make_chunk(
.to_string()
}
fn make_usage_chunk(
id: &str,
model: &str,
created: u64,
prompt_tokens: usize,
completion_tokens: usize,
) -> String {
serde_json::json!({
"id": id,
"object": "chat.completion.chunk",
"created": created,
"model": model,
"choices": [],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
}
})
.to_string()
}
fn unix_timestamp() -> u64 {
std::time::SystemTime::now()
.duration_since(std::time::UNIX_EPOCH)
@@ -558,9 +692,20 @@ fn sampling_params(req: &ChatRequest) -> SamplingParams {
temperature: req.temperature.unwrap_or(0.0),
top_k: req.top_k.unwrap_or(0),
top_p: req.top_p.unwrap_or(1.0),
presence_penalty: req.presence_penalty.unwrap_or(0.0),
repetition_penalty: req.repetition_penalty.unwrap_or(1.0),
}
}
fn template_options(req: &ChatRequest) -> (Option<bool>, Option<bool>) {
let kwargs = req.chat_template_kwargs.as_ref();
(
req.enable_thinking
.or_else(|| kwargs.and_then(|value| value.enable_thinking)),
kwargs.and_then(|value| value.preserve_thinking),
)
}
/// Map engine finish_reason strings to OpenAI-standard values. Any engine-internal
/// code (e.g. "error" from tp/pp client-stall) collapses to None so SDK clients see
/// a clean null instead of an unknown value.

21
crates/xserv-server/src/engine.rs
View File

@@ -4,7 +4,9 @@ use std::sync::Once;
use std::sync::mpsc;
use std::time::Instant;
use xserv_model::loader;
use xserv_model::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample};
use xserv_model::{
BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, SamplingParams, sample_with_history,
};
use xserv_tensor::{DType, Device};
use xserv_tokenizer::Tokenizer;
@@ -232,7 +234,7 @@ impl Engine {
slot,
&mut self.paged_cache,
);
let next = sample(&logits, &seq.sampling);
let next = sample_with_history(&logits, &seq.sampling, &seq.prompt_tokens);
seq.generated_tokens.push(next);
seq.prefilled = true;
emit_token(&self.tokenizer, seq, next);
@@ -310,7 +312,11 @@ impl Engine {
// (~1.2 MB for B=4, Qwen3 vocab=152K).
let all_greedy = decode_indices
.iter()
.all(|&i| running[i].sampling.temperature == 0.0);
.all(|&i| {
running[i].sampling.temperature == 0.0
&& running[i].sampling.presence_penalty == 0.0
&& running[i].sampling.repetition_penalty == 1.0
});
if all_greedy {
let next_ids = xserv_kernels::argmax_bf16_to_host(&logits);
for (j, &i) in decode_indices.iter().enumerate() {
@@ -329,7 +335,10 @@ impl Engine {
for (j, &i) in decode_indices.iter().enumerate() {
let row_start = j * vocab_size;
let row_logits = &data[row_start..row_start + vocab_size];
let next = if running[i].sampling.temperature == 0.0 {
let unpenalized_greedy = running[i].sampling.temperature == 0.0
&& running[i].sampling.presence_penalty == 0.0
&& running[i].sampling.repetition_penalty == 1.0;
let next = if unpenalized_greedy {
row_logits
.iter()
.enumerate()
@@ -339,7 +348,9 @@ impl Engine {
} else {
let row_tensor =
xserv_tensor::Tensor::from_slice(row_logits, &[1, vocab_size]);
sample(&row_tensor, &running[i].sampling)
let mut history = running[i].prompt_tokens.clone();
history.extend_from_slice(&running[i].generated_tokens);
sample_with_history(&row_tensor, &running[i].sampling, &history)
};
running[i].generated_tokens.push(next);
emit_token(&self.tokenizer, &mut running[i], next);

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

@@ -10,8 +10,9 @@ use axum::{
};
use engine::GenerateRequest;
use std::path::PathBuf;
use std::sync::atomic::AtomicBool;
use std::sync::{Arc, Mutex, mpsc};
use xserv_model::ModelConfig;
use xserv_model::{ModelConfig, ModelFamily};
pub struct AppState {
pub model_name: String,
@@ -19,6 +20,7 @@ pub struct AppState {
pub engine_sender: Mutex<mpsc::SyncSender<GenerateRequest>>,
pub engine_tokenizer: Mutex<xserv_tokenizer::Tokenizer>,
pub max_seq_len: usize,
pub engine_ready: Arc<AtomicBool>,
}
#[tokio::main]
@@ -77,12 +79,18 @@ async fn main() {
std::process::exit(1);
}
let model_config = ModelConfig::from_file(&model_dir.join("config.json"));
// gpt-oss is only implemented in the TP engine; route it there even at
// tp=1 (single-rank world) so quantized models can serve on one GPU.
let is_gpt_oss = model_config.model_type.as_deref() == Some("gpt_oss");
if pp > 1 && is_gpt_oss {
let model_family = model_config.family();
if model_family == ModelFamily::Qwen35Moe && pp <= 1 {
eprintln!(
"gpt-oss is not supported by the pipeline-parallel engine (Qwen3 only); use --tp instead"
"Qwen3.5/3.6 MoE model '{}' currently requires pipeline parallelism; use --pp 8 on dash5",
model_config.model_type_str()
);
std::process::exit(1);
}
if pp > 1 && !matches!(model_family, ModelFamily::Qwen3 | ModelFamily::Qwen35Moe) {
eprintln!(
"{} is not supported by the pipeline-parallel engine; use --tp instead",
model_family.as_str()
);
std::process::exit(1);
}
@@ -109,27 +117,42 @@ async fn main() {
// behind, instead of letting them pile up in RAM. try_send in the API
// handler surfaces this as 503 to the client.
let (tx, rx) = mpsc::sync_channel::<GenerateRequest>(256);
let engine_ready = Arc::new(AtomicBool::new(false));
let model_dir_clone = model_dir.clone();
let engine_ready_clone = Arc::clone(&engine_ready);
std::thread::spawn(move || {
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 && !is_gpt_oss {
pp_engine::run_pp(
&model_dir_clone,
pp,
max_seq_len,
rx,
engine_ready_clone,
);
} else if tp <= 1 && model_family == ModelFamily::Qwen3 {
let mut engine = engine::Engine::load_with_swap(
&model_dir_clone,
max_batch,
max_seq_len,
swap_space_gb,
);
engine_ready_clone.store(true, std::sync::atomic::Ordering::Release);
engine.run(rx);
} else {
// Tensor-parallel path: rank-0 coordinator + worker rank threads.
tp_engine::run_tp(&model_dir_clone, tp, max_seq_len, rx);
tp_engine::run_tp(
&model_dir_clone,
tp,
max_seq_len,
rx,
engine_ready_clone,
);
}
});
let model_type = model_config.model_type.clone().unwrap_or_default();
let model_type = model_config.model_type_str().to_string();
let chat_template = api::ChatTemplate::load(&model_dir, &model_type);
let state = Arc::new(AppState {
model_name,
@@ -137,6 +160,7 @@ async fn main() {
engine_sender: Mutex::new(tx),
engine_tokenizer: Mutex::new(tokenizer),
max_seq_len,
engine_ready,
});
let app = Router::new()

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

@@ -16,14 +16,18 @@
use std::ffi::c_void;
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
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::{BLOCK_SIZE, ModelConfig, PagedKVCache, Qwen3, sample};
use xserv_model::sampling::{SamplingParams, sample_with_history};
use xserv_model::{
BLOCK_SIZE, ModelConfig, ModelFamily, PagedKVCache, Qwen3, Qwen35Moe,
Qwen35RecurrentCache,
};
use xserv_tensor::{DType, Device, Tensor};
use xserv_tokenizer::Tokenizer;
@@ -39,7 +43,7 @@ enum PpCommand {
/// 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,
tokens: Vec<u32>,
slot: usize,
sampling: SamplingParams,
},
@@ -52,13 +56,69 @@ enum PpCommand {
}
struct StageCtx {
model: Qwen3,
model: PpModel,
qwen35_recurrent: Option<Qwen35RecurrentCache>,
cache: PagedKVCache,
pp: Arc<PpContext>,
hidden: usize,
device: u32,
}
enum PpModel {
Qwen3(Qwen3),
Qwen35(Qwen35Moe),
}
impl StageCtx {
fn reset_slot(&mut self, slot: usize) {
if let Some(cache) = &mut self.qwen35_recurrent {
cache.reset_slot(slot);
}
}
fn embed(&self, tokens: &[u32]) -> Tensor {
match &self.model {
PpModel::Qwen3(model) => model.embed(tokens),
PpModel::Qwen35(model) => model.embed(tokens),
}
}
fn head(&self, x: &Tensor) -> Tensor {
match &self.model {
PpModel::Qwen3(model) => model.head(x),
PpModel::Qwen35(model) => model.head(x),
}
}
fn forward_prefill(&mut self, x: Tensor, slot: usize) -> Tensor {
match &self.model {
PpModel::Qwen3(model) => model.forward_layers_prefill(x, slot, &mut self.cache),
PpModel::Qwen35(model) => model.forward_layers_prefill(
x,
slot,
&mut self.cache,
self.qwen35_recurrent
.as_mut()
.expect("Qwen3.6 recurrent cache"),
),
}
}
fn forward_decode(&mut self, x: Tensor, slot: usize) -> Tensor {
match &self.model {
PpModel::Qwen3(model) => model.forward_layers_decode(x, &[slot], &mut self.cache),
PpModel::Qwen35(model) => model.forward_layers_decode(
x,
slot,
&mut self.cache,
self.qwen35_recurrent
.as_mut()
.expect("Qwen3.6 recurrent cache"),
),
}
}
}
/// Build this stage: NCCL init, load + slice weights, size a per-stage KV pool
/// for THIS stage's layers only (so per-GPU KV is ~1/P).
fn build_stage(
@@ -71,14 +131,40 @@ fn build_stage(
id: UniqueId,
) -> StageCtx {
let pp = Arc::new(PpContext::init(stage, world, id, device));
let weights = loader::load_model_dir(model_dir, Device::Cpu);
let model = Qwen3::from_weights_pp(config.clone(), weights, stage, world, device);
let model = match config.family() {
ModelFamily::Qwen35Moe => {
let weights = loader::load_model_dir_filtered(model_dir, Device::Cpu, |name| {
Qwen35Moe::weight_belongs_to_pp_stage(config, name, stage, world)
});
PpModel::Qwen35(Qwen35Moe::from_weights_pp(
config.clone(),
weights,
stage,
world,
device,
))
}
_ => {
let weights = loader::load_model_dir(model_dir, Device::Cpu);
PpModel::Qwen3(Qwen3::from_weights_pp(
config.clone(),
weights,
stage,
world,
device,
))
}
};
// The KV cache only needs this stage's layers; build it from a config clone
// whose layer count is the per-stage count (heads are NOT split under PP).
let per_stage = config.num_layers() / world;
let mut stage_config = config.clone();
stage_config.num_hidden_layers = Some(per_stage);
if let Some(text) = stage_config.text_config.as_mut() {
text.num_hidden_layers = Some(per_stage);
} else {
stage_config.num_hidden_layers = Some(per_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
@@ -91,8 +177,13 @@ fn build_stage(
DType::BF16,
device,
);
let qwen35_recurrent = match &model {
PpModel::Qwen35(model) => Some(model.new_recurrent_cache(4)),
PpModel::Qwen3(_) => None,
};
StageCtx {
model,
qwen35_recurrent,
cache,
pp,
hidden: config.hidden(),
@@ -138,40 +229,51 @@ fn worker_loop(
max_seq_len,
id,
);
let _ = ack_tx.send(());
let is_last = stage == world - 1;
let mut token_history = vec![Vec::<u32>::new(); 4];
let prev = stage - 1;
let next = stage + 1;
while let Ok(cmd) = cmd_rx.recv() {
match cmd {
PpCommand::Register(slot) => {
token_history[slot].clear();
sc.reset_slot(slot);
let _ = sc.cache.register_sequence(slot);
let _ = ack_tx.send(());
}
PpCommand::Free(slot) => {
token_history[slot].clear();
sc.cache.free_sequence(slot);
let _ = ack_tx.send(());
}
PpCommand::Prefill {
n_tokens,
tokens,
slot,
sampling,
} => {
let n_tokens = tokens.len();
token_history[slot] = tokens;
let x = recv_hidden(&sc, n_tokens, prev);
let x = sc.model.forward_layers_prefill(x, slot, &mut sc.cache);
let x = sc.forward_prefill(x, slot);
if is_last {
let logits = sc.model.head(&x);
let _ = token_tx.send(sample(&logits, &sampling));
let logits = sc.head(&x);
let token = sample_with_history(&logits, &sampling, &token_history[slot]);
token_history[slot].push(token);
let _ = token_tx.send(token);
} else {
send_hidden(&sc, &x, next);
}
}
PpCommand::Decode { slot, sampling } => {
let x = recv_hidden(&sc, 1, prev);
let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache);
let x = sc.forward_decode(x, slot);
if is_last {
let logits = sc.model.head(&x);
let _ = token_tx.send(sample(&logits, &sampling));
let logits = sc.head(&x);
let token = sample_with_history(&logits, &sampling, &token_history[slot]);
token_history[slot].push(token);
let _ = token_tx.send(token);
} else {
send_hidden(&sc, &x, next);
}
@@ -191,6 +293,7 @@ pub fn run_pp(
world: usize,
max_seq_len: usize,
rx: mpsc::Receiver<GenerateRequest>,
ready: Arc<AtomicBool>,
) {
assert!(world >= 2, "run_pp requires world >= 2");
let config = ModelConfig::from_file(&model_dir.join("config.json"));
@@ -231,6 +334,10 @@ pub fn run_pp(
// Stage 0 (this thread): coordinator + embedding + first layers.
let mut sc = build_stage(model_dir, &config, 0, world, 0, max_seq_len, id);
for _ in 1..world {
ack_rx.recv().expect("PP worker exited during model load");
}
ready.store(true, Ordering::Release);
eprintln!("[pp-engine] ready (pp={world}, max_seq_len={max_seq_len})");
let n_workers = world - 1;
@@ -249,6 +356,7 @@ pub fn run_pp(
let slot = 0usize;
while let Ok(req) = rx.recv() {
broadcast(&cmd_txs, PpCommand::Register(slot));
sc.reset_slot(slot);
sc.cache.register_sequence(slot).expect("register slot");
wait_acks(&ack_rx);
@@ -256,13 +364,13 @@ pub fn run_pp(
broadcast(
&cmd_txs,
PpCommand::Prefill {
n_tokens: req.prompt_tokens.len(),
tokens: req.prompt_tokens.clone(),
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);
let x = sc.embed(&req.prompt_tokens);
let x = sc.forward_prefill(x, slot);
send_hidden(&sc, &x, next_peer);
let mut next = token_rx.recv().expect("prefill token");
@@ -287,8 +395,8 @@ pub fn run_pp(
sampling: req.sampling.clone(),
},
);
let x = sc.model.embed(&[next]);
let x = sc.model.forward_layers_decode(x, &[slot], &mut sc.cache);
let x = sc.embed(&[next]);
let x = sc.forward_decode(x, slot);
send_hidden(&sc, &x, next_peer);
next = token_rx.recv().expect("decode token");
generated += 1;

37
crates/xserv-server/src/tp_engine.rs
View File

@@ -14,14 +14,15 @@
use std::path::{Path, PathBuf};
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::thread;
use xserv_distributed::{TpContext, UniqueId};
use xserv_model::loader;
use xserv_model::{
BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, PagedKVCache, Qwen3, sample,
sample_greedy_penalized,
BLOCK_SIZE, GptOss, GraphedGptOssDecoder, ModelConfig, ModelFamily, PagedKVCache, Qwen3,
sample_greedy_penalized, sample_with_history,
};
use xserv_tensor::{DType, Device, Tensor};
use xserv_tokenizer::Tokenizer;
@@ -105,24 +106,26 @@ fn build_rank(
tp: Option<Arc<TpContext>>,
) -> RankCtx {
let weights = loader::load_model_dir(model_dir, Device::Cpu);
let model = if config.is_moe() {
TpModel::GptOss(GptOss::from_weights_tp(
let model = match config.family() {
ModelFamily::GptOss => TpModel::GptOss(GptOss::from_weights_tp(
config.clone(),
weights,
rank,
world,
device,
tp,
))
} else {
TpModel::Qwen3(Qwen3::from_weights_tp(
)),
ModelFamily::Qwen35Moe => {
panic!("Qwen3.5/3.6 MoE reached TP engine before inference support was implemented")
}
_ => TpModel::Qwen3(Qwen3::from_weights_tp(
config.clone(),
weights,
rank,
world,
device,
tp,
))
)),
};
let local_kv = config.num_kv_heads() / world;
let max_blocks_per_seq = (max_seq_len + BLOCK_SIZE - 1) / BLOCK_SIZE;
@@ -164,6 +167,7 @@ fn worker_loop(
max_seq_len,
Some(tp),
);
let _ = ack_tx.send(());
while let Ok(cmd) = cmd_rx.recv() {
match cmd {
TpCommand::Register(slot) => {
@@ -196,6 +200,7 @@ pub fn run_tp(
world: usize,
max_seq_len: usize,
rx: mpsc::Receiver<GenerateRequest>,
ready: Arc<AtomicBool>,
) {
// world=1 is a valid single-rank configuration (gpt-oss has no
// single-GPU engine path; NCCL init and all_reduce no-op at world=1).
@@ -235,6 +240,10 @@ pub fn run_tp(
// Rank 0 (this thread).
let tp = Arc::new(TpContext::init(0, world, id, 0));
let mut rc = build_rank(model_dir, &config, 0, world, 0, max_seq_len, Some(tp));
for _ in 1..world {
ack_rx.recv().expect("TP worker exited during model load");
}
ready.store(true, Ordering::Release);
eprintln!("[tp-engine] ready (tp={world}, max_seq_len={max_seq_len})");
// Optional repetition penalty to break greedy repetition loops (reasoning
@@ -254,7 +263,7 @@ pub fn run_tp(
let start = history.len().saturating_sub(rep_window);
sample_greedy_penalized(logits, &history[start..], rep_penalty)
} else {
sample(logits, sp)
sample_with_history(logits, sp, history)
}
};
@@ -288,9 +297,9 @@ pub fn run_tp(
.model
.forward_prefill_paged(&req.prompt_tokens, slot, &mut rc.cache);
wait_acks(&ack_rx);
let mut gen_ids: Vec<u32> = Vec::new();
let mut next = pick(&logits, &req.sampling, &gen_ids);
gen_ids.push(next);
let mut token_history = req.prompt_tokens.clone();
let mut next = pick(&logits, &req.sampling, &token_history);
token_history.push(next);
let mut decode_buf: Vec<u8> = Vec::new();
let mut generated = 1usize;
@@ -317,8 +326,8 @@ pub fn run_tp(
);
let logits = rank_decode(&mut rc, &[next], &[pos], &[slot]);
wait_acks(&ack_rx);
next = pick(&logits, &req.sampling, &gen_ids);
gen_ids.push(next);
next = pick(&logits, &req.sampling, &token_history);
token_history.push(next);
generated += 1;
stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
};