Skip to content

perf(frontend): skip per-token text decode for chunks without logprobs #17349

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jthomson04 merged 3 commits into from
Oct 27, 2026
Merged

perf(frontend): skip per-token text decode for chunks without logprobs #17349

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
ebb113e
perf(frontend): skip per-token text decode for chunks without logprobs
jthomson04 Oct 26, 2026
ac62715
perf(frontend): drop discarded per-token text clones in Decoder::step
jthomson04 Oct 26, 2026
395d239
refactor(frontend): StepResult constructors no longer take a token
jthomson04 Oct 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
204 changes: 116 additions & 88 deletions lib/llm/src/backend.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -454,7 +454,9 @@ impl
return Some((output, state));
};

let result = decoder.process_token_ids(&data.token_ids).and_then(|mut result| {
// Per-token text is read only with this chunk's logprobs (`create_logprobs`).
let needs_token_text = data.log_probs.is_some();
let result = decoder.process_token_ids_with_token_text(&data.token_ids, needs_token_text).and_then(|mut result| {
if result.stop_trigger.is_none() && data.finish_reason.is_some() {
let final_step = decoder.finish()?;
if let Some(text) = final_step.released_text {
Expand Down Expand Up @@ -770,7 +772,8 @@ pub struct StepResult {
/// `SeqResult.tokens[i]` must equal the text of `token_ids[i]` so that consumers
/// zipping `tokens` with per-token logprobs (e.g. `create_logprobs`) stay aligned --
/// logprobs describe what the model generated and must not change because a hidden
/// stop sequence also happens to be in flight.
/// stop sequence also happens to be in flight. `None` when the caller did not ask for
/// per-token text (see `Decoder::step_with_token_text`).
pub token: Option<String>,
/// Text to append to the caller-visible `text`/content this step. May be `None` on a
/// step whose text is being withheld as a possible hidden-stop-sequence prefix, and may
Expand All @@ -781,33 +784,17 @@ pub struct StepResult {
}

impl StepResult {
/// No stop-sequence withholding in effect: the token's own text is released as-is.
fn ok(token: Option<String>) -> Self {
fn ok(released_text: Option<String>) -> Self {
Self {
token: token.clone(),
released_text: token,
stop_trigger: None,
}
}

/// `token` and `released_text` differ, e.g. because `released_text` also carries a
/// previously withheld backlog, or because a matched stop sequence is excluded from
/// `released_text` but not from `token`.
fn ok_split(token: Option<String>, released_text: Option<String>) -> Self {
Self {
token,
token: None,
released_text,
stop_trigger: None,
}
}

fn with_stop_trigger(
token: Option<String>,
released_text: Option<String>,
stop_trigger: StopTrigger,
) -> Self {
fn with_stop_trigger(released_text: Option<String>, stop_trigger: StopTrigger) -> Self {
Self {
token,
token: None,
released_text,
stop_trigger: Some(stop_trigger),
}
Expand Down Expand Up @@ -983,6 +970,16 @@ impl Decoder {
/// In the future, this method may kick off async cpu/tokio tasks and or async cuda tasks to
/// handle logits post-processing and/or other tasks.
pub fn step(&mut self, token_id: TokenIdType) -> Result<StepResult> {
self.step_with_token_text(token_id, true)
}

/// Like [`Self::step`], but leaves `StepResult::token` as `None` when `needs_token_text`
/// is false.
fn step_with_token_text(
&mut self,
token_id: TokenIdType,
needs_token_text: bool,
) -> Result<StepResult> {
let below_min_tokens = self.generated_tokens < self.min_tokens;
// increment the generated tokens
self.generated_tokens += 1;
Expand All @@ -994,7 +991,6 @@ impl Decoder {
&& !self.no_stop_trim
{
let mut result = self.finish()?;
result.token = None;
result.stop_trigger.get_or_insert_with(|| {
if self.user_stop_ids.contains(&token_id) {
StopTrigger::UserStopTokenDetected(token_id)
Expand All @@ -1005,14 +1001,6 @@ impl Decoder {
return Ok(result);
}

// Decode the selected token independently for per-token metadata. Incremental
// decoding may release a backlog of earlier byte-fallback IDs on this step; using
// that backlog as this ID's token string would attach it to the wrong logprob.
let selected_token = self
.tokenizer
.as_ref()
.and_then(|tokenizer| decode_token_text(tokenizer, token_id, self.skip_special_tokens));

// Decode caller-visible content. Track IDs retained by the incremental decoder so
// a migration checkpoint can replay them instead of losing buffered byte fallback.
let detokenize_start = self.tracker.as_ref().map(|_| Instant::now());
Expand All @@ -1027,11 +1015,27 @@ impl Decoder {
if let (Some(start), Some(tracker)) = (detokenize_start, &self.tracker) {
tracker.record_detokenize_latency(start.elapsed());
}
let selected_token = selected_token.or_else(|| decoded_text.clone());

// Decode the selected token independently for per-token metadata. Incremental
// decoding may release a backlog of earlier byte-fallback IDs on this step; using
// that backlog as this ID's token string would attach it to the wrong logprob.
let selected_token = if needs_token_text {
self.tokenizer
.as_ref()
.and_then(|tokenizer| {
decode_token_text(tokenizer, token_id, self.skip_special_tokens)
})
.or_else(|| decoded_text.clone())
} else {
None
};

// stop conditions to not apply until the minimum number of tokens have been generated
if below_min_tokens {
return Ok(StepResult::ok_split(selected_token, decoded_text));
return Ok(StepResult {
token: selected_token,
..StepResult::ok(decoded_text)
});
}

if self.visible_stop_ids.contains(&token_id) || self.hidden_stop_ids.contains(&token_id) {
Expand All @@ -1045,57 +1049,33 @@ impl Decoder {
// Stop strings are evaluated before stop-token IDs so a configured string ending
// at a visible stop token cannot leak through the token-stop fast path.
let mut decoded = self.process_decoded_text(decoded_text)?;
decoded.token = selected_token.clone();
decoded.token = selected_token;
if decoded.stop_trigger.is_some() {
return Ok(decoded);
}

// Check token stops. Visible token IDs are included in output.
if self.visible_stop_ids.contains(&token_id) {
// Any text still withheld as a partial hidden-stop-sequence match can never
// complete now, so it goes out ahead of this (included) token's own text --
// otherwise it would be silently dropped and, if it were released later, would
// come out of order.
if let Some(flushed) = self.flush_jailed() {
decoded
.released_text
.get_or_insert_default()
.push_str(&flushed);
}
return Ok(StepResult::with_stop_trigger(
selected_token,
decoded.released_text,
StopTrigger::VisibleStopTokenDetected(token_id),
));
}

// Check token stops. User-provided token IDs take precedence over
// system/EOS IDs so stop_reason only reports stops the caller requested.
if self.hidden_stop_ids.contains(&token_id) {
let trigger = if self.user_stop_ids.contains(&token_id) {
StopTrigger::UserStopTokenDetected(token_id)
} else {
StopTrigger::HiddenStopTokenDetected(token_id)
};
// Release any earlier withheld prefix before this stop token's own text.
let token = if self.no_stop_trim {
selected_token
} else {
None
};
if let Some(flushed) = self.flush_jailed() {
decoded
.released_text
.get_or_insert_default()
.push_str(&flushed);
}
return Ok(StepResult::with_stop_trigger(
token,
decoded.released_text,
trigger,
));
// Check token stops. Visible token IDs are included in output. User-provided token
// IDs take precedence over system/EOS IDs so stop_reason only reports stops the
// caller requested. A hidden ID only reaches here under no_stop_trim.
let trigger = if self.visible_stop_ids.contains(&token_id) {
StopTrigger::VisibleStopTokenDetected(token_id)
} else if self.user_stop_ids.contains(&token_id) {
StopTrigger::UserStopTokenDetected(token_id)
} else if self.hidden_stop_ids.contains(&token_id) {
StopTrigger::HiddenStopTokenDetected(token_id)
} else {
return Ok(decoded);
};
// Any text still withheld as a partial hidden-stop-sequence match can never
// complete now, so it goes out ahead of this token's own text -- otherwise it would
// be silently dropped and, if it were released later, would come out of order.
if let Some(flushed) = self.flush_jailed() {
decoded
.released_text
.get_or_insert_default()
.push_str(&flushed);
}

decoded.stop_trigger = Some(trigger);
Ok(decoded)
}

Expand Down Expand Up @@ -1124,11 +1104,7 @@ impl Decoder {
.then(|| self.jail[release_start..offset].to_string())
.filter(|s| !s.is_empty());
self.jailed_bytes = 0;
// `decoded_text` remains provisional metadata here; `step` replaces it
// with the independently decoded selected token before returning. Only
// caller-visible `released_text` excludes the matched sequence.
return Ok(StepResult::with_stop_trigger(
decoded_text,
partial_token,
StopTrigger::HiddenStopSequenceDetected(seq.to_string()),
));
Expand All @@ -1147,7 +1123,6 @@ impl Decoder {
.filter(|s| !s.is_empty());
self.jailed_bytes = 0;
return Ok(StepResult::with_stop_trigger(
decoded_text,
token_with_stop,
StopTrigger::VisibleStopSequenceDetected(seq.to_string()),
));
Expand All @@ -1164,9 +1139,7 @@ impl Decoder {
.then(|| self.jail[release_start..release_end].to_string());

Self::maybe_drain_to_max_bytes(&mut self.jail, self.jail_max_bytes);
// `decoded_text` is provisional per-token metadata that `step` replaces with
// the selected token; `released` is what newly clears the jail for content.
return Ok(StepResult::ok_split(decoded_text, released));
return Ok(StepResult::ok(released));
}

Ok(StepResult::ok(decoded_text))
Expand Down Expand Up @@ -1279,6 +1252,16 @@ impl Decoder {
}

pub fn process_token_ids(&mut self, token_ids: &[TokenIdType]) -> Result<SeqResult> {
self.process_token_ids_with_token_text(token_ids, true)
}

/// Like [`Self::process_token_ids`], but every `tokens[i]` is `None` when
/// `needs_token_text` is false. `tokens` still has one entry per consumed ID.
fn process_token_ids_with_token_text(
&mut self,
token_ids: &[TokenIdType],
needs_token_text: bool,
) -> Result<SeqResult> {
let mut text: Option<String> = None;
let mut tokens = Vec::with_capacity(token_ids.len());

Expand All @@ -1287,7 +1270,7 @@ impl Decoder {
token,
released_text,
stop_trigger,
} = self.step(*token_id)?;
} = self.step_with_token_text(*token_id, needs_token_text)?;

// `text` accumulates the caller-visible content, which can lag behind and later
// release more than one step's worth at once. `tokens[i]` always reports
Expand Down Expand Up @@ -1465,8 +1448,10 @@ mod tests {
stop: stop.map(|text| vec![text.to_string()]),
..Default::default()
};
// Per-token text is only decoded for chunks that carry logprobs.
let mut outputs = vec![LLMEngineOutput {
token_ids: vec![0],
log_probs: Some(vec![0.0]),
..Default::default()
}];
match boundary {
Expand Down Expand Up @@ -1844,6 +1829,49 @@ mod tests {
}
}

#[tokio::test]
async fn test_token_text_only_decoded_for_chunks_with_logprobs() {
let tokenizer: Arc<dyn traits::Tokenizer> = Arc::new(CandidateDecoder);
let backend = Backend::from_tokenizer(Tokenizer::from(tokenizer));
let request = PreprocessedRequest::builder()
.model("test-model".to_string())
.token_ids(vec![])
.stop_conditions(StopConditions::default())
.sampling_options(SamplingOptions::default())
.output_options(OutputOptions::default())
.build()
.unwrap();
let engine: ServerStreamingEngine<PreprocessedRequest, Annotated<LLMEngineOutput>> =
Arc::new(SyntheticSglangStopEngine {
outputs: Some(vec![
LLMEngineOutput {
token_ids: vec![101],
log_probs: Some(vec![-0.1]),
..Default::default()
},
LLMEngineOutput {
token_ids: vec![103],
..Default::default()
},
]),
});
let outputs: Vec<_> = Operator::generate(backend.as_ref(), SingleIn::new(request), engine)
.await
.unwrap()
.filter_map(|item| async move { item.data })
.map(|data| (data.text, data.tokens))
.collect()
.await;

assert_eq!(
outputs,
vec![
(Some("Okay".to_string()), vec![Some("Okay".to_string())]),
(Some("<|user|>".to_string()), vec![None]),
]
);
}

#[tokio::test]
async fn test_sglang_local_stop_discards_all_suffix_tokens() {
use crate::protocols::common::extensions::NvExt;
Expand Down
Loading