server: non-blocking stream send — stop one slow client stalling the batch

All three engines emitted tokens with blocking_send on the single
decode/coordinator OS thread. A streaming client that drains slower than
generation fills its 64-slot channel, and blocking_send then blocks the whole
thread: under continuous batching one slow consumer stalls every other running
sequence (and in the serial TP/PP path it blocks admission of the next request
too). The whole point of continuous batching is defeated.

Fix: switch to try_send. engine.rs sets a client_stalled flag on Full/Closed,
reaped by is_finished() next iteration; tp_engine/pp_engine emit_text returns
bool and the decode loop breaks with finish_reason "error". When the
sequence/request is dropped its sender drops too, closing the channel so the
client receive loop ends rather than hanging. A slow client now only loses its
own sequence, never the batch.

Verified on dash5: gpt-oss FP8 TP=1 streaming via tp_engine still streams
correctly (SSE chunks, coherent content, no hang); bench-gpt-oss TP=2 5.9ms
TPOT unchanged.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
2026-07-01 12:37:32 +08:00
parent cfbd64d206
commit 0314b4f3ac
3 changed files with 60 additions and 20 deletions

View File

@@ -294,9 +294,12 @@ pub fn run_tp(
let mut decode_buf: Vec<u8> = Vec::new();
let mut generated = 1usize;
emit_text(&tokenizer, &req, next, &mut decode_buf);
let mut stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
let finish = loop {
if stalled {
break "error";
}
if tokenizer.is_eos(next) {
break "stop";
}
@@ -317,17 +320,17 @@ pub fn run_tp(
next = pick(&logits, &req.sampling, &gen_ids);
gen_ids.push(next);
generated += 1;
emit_text(&tokenizer, &req, next, &mut decode_buf);
stalled = !emit_text(&tokenizer, &req, next, &mut decode_buf);
};
let tail = tokenizer.flush_decode_stream(&mut decode_buf);
if !tail.is_empty() {
let _ = req.sender.blocking_send(GenerateEvent::Token {
let _ = req.sender.try_send(GenerateEvent::Token {
id: next,
text: tail,
});
}
let _ = req.sender.blocking_send(GenerateEvent::Done {
let _ = req.sender.try_send(GenerateEvent::Done {
finish_reason: finish.to_string(),
});
@@ -340,14 +343,19 @@ pub fn run_tp(
}
/// Stream a token's decoded text to the client (EOS contributes no text).
fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &mut Vec<u8>) {
/// Returns false if the send would block (client too slow) or the client is
/// gone — the caller stops generating so the serial coordinator thread is free
/// to admit the next request instead of blocking on one slow consumer.
fn emit_text(tokenizer: &Tokenizer, req: &GenerateRequest, token_id: u32, buf: &mut Vec<u8>) -> bool {
if tokenizer.is_eos(token_id) {
return;
return true;
}
let text = tokenizer.decode_token_stream(token_id, buf);
if !text.is_empty() {
let _ = req
return req
.sender
.blocking_send(GenerateEvent::Token { id: token_id, text });
.try_send(GenerateEvent::Token { id: token_id, text })
.is_ok();
}
true
}