diff --git a/crates/tokenizer/src/mock.rs b/crates/tokenizer/src/mock.rs index f057aabe2..245e1968b 100644 --- a/crates/tokenizer/src/mock.rs +++ b/crates/tokenizer/src/mock.rs @@ -122,6 +122,11 @@ impl TokenizerTrait for MockTokenizer { self.reverse_vocab.get(&id).cloned() } + fn eos_token_ids(&self) -> &[u32] { + // `` in the mock vocab. + &[999] + } + fn as_any(&self) -> &dyn std::any::Any { self } diff --git a/crates/tokenizer/src/stop.rs b/crates/tokenizer/src/stop.rs index 820620cec..3b0ed3e6f 100644 --- a/crates/tokenizer/src/stop.rs +++ b/crates/tokenizer/src/stop.rs @@ -77,6 +77,9 @@ pub struct StopSequenceDecoder { jail_max_bytes: usize, /// Whether we've stopped stopped: bool, + /// The string stop sequence that triggered the stop, if any. Set only for + /// string-sequence matches; token-level stops leave this `None`. + matched_stop: Option, /// True when there are no string stop sequences (only token-level stops). /// In this mode the jail buffer is bypassed entirely for lower overhead. token_only: bool, @@ -139,6 +142,7 @@ impl StopSequenceDecoder { jail_buffer: String::new(), jail_max_bytes, stopped: false, + matched_stop: None, token_only, } } @@ -208,6 +212,7 @@ impl StopSequenceDecoder { let input = Input::new(&self.jail_buffer).span(search_start..self.jail_buffer.len()); if let Some(mat) = ac.find(input) { self.stopped = true; + self.matched_stop = Some(self.jail_buffer[mat.start()..mat.end()].to_string()); let is_visible = mat.pattern().as_usize() >= self.visible_boundary_idx; if is_visible { @@ -287,11 +292,18 @@ impl StopSequenceDecoder { self.stopped } + /// The string stop sequence that triggered the stop, if a string sequence + /// matched. `None` for token-level stops or when no stop has fired. + pub fn matched_stop(&self) -> Option<&str> { + self.matched_stop.as_deref() + } + /// Reset the decoder state pub fn reset(&mut self) { self.jail_buffer.clear(); self.sequence.clear(); self.stopped = false; + self.matched_stop = None; } } @@ -455,6 +467,40 @@ mod tests { )); } + #[test] + fn test_matched_stop_reports_matched_string() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_sequence("test"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // No match yet. + assert_eq!(decoder.matched_stop(), None); + + decoder.process_token(1).unwrap(); // "Hello" + decoder.process_token(2).unwrap(); // "world" + assert_eq!(decoder.matched_stop(), None); + + // "test" triggers the string stop; the matched string is captured. + decoder.process_token(3).unwrap(); // "test" + assert_eq!(decoder.matched_stop(), Some("test")); + + // Reset clears it. + decoder.reset(); + assert_eq!(decoder.matched_stop(), None); + } + + #[test] + fn test_matched_stop_none_for_token_stop() { + // A token-id stop is not a string match, so matched_stop stays None. + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_token(999); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + let result = decoder.process_token(999).unwrap(); + assert_eq!(result, SequenceDecoderOutput::Stopped); + assert_eq!(decoder.matched_stop(), None); + } + #[test] fn test_flush_after_partial() { let tokenizer = Arc::new(MockTokenizer::new()); diff --git a/model_gateway/src/routers/grpc/backend_client.rs b/model_gateway/src/routers/grpc/backend_client.rs index f7ede0bba..e5fa7e97a 100644 --- a/model_gateway/src/routers/grpc/backend_client.rs +++ b/model_gateway/src/routers/grpc/backend_client.rs @@ -49,12 +49,10 @@ impl BackendClient { } } - /// True if this backend speaks the vLLM protocol (gRPC-vLLM or ZMQ). - pub fn is_vllm(&self) -> bool { - match self { - Self::Grpc(client) => client.is_vllm(), - Self::Zmq(_) => true, - } + /// True if this is a direct-ZMQ backend (the engine receives token ids only + /// and cannot match string stops itself). + pub fn is_zmq(&self) -> bool { + matches!(self, Self::Zmq(_)) } /// Local liveness. gRPC has no cheap local flag (it uses a health RPC), so diff --git a/model_gateway/src/routers/grpc/common/stages/helpers.rs b/model_gateway/src/routers/grpc/common/stages/helpers.rs index 3cecfe4a6..fd3fbdb55 100644 --- a/model_gateway/src/routers/grpc/common/stages/helpers.rs +++ b/model_gateway/src/routers/grpc/common/stages/helpers.rs @@ -2,6 +2,7 @@ use std::sync::Arc; +use llm_tokenizer::traits::Tokenizer; use rand::RngExt; use smg_grpc_client::{ mlx_proto, @@ -291,6 +292,130 @@ fn apply_tokenspeed_sampling_defaults( apply_opt!(repetition_penalty); } +/// Convert single-token stop strings into `stop_token_ids` entries so the engine +/// can halt generation early for the common case (e.g. `["."]`, `["\n"]`). +/// +/// The proto `stop_token_ids` field is a flat list of single token ids, so a +/// multi-token stop string cannot be represented there — pushing its sub-tokens +/// would stop far too eagerly (on any one of them). Multi-token, empty, and +/// unknown stops are therefore left to the router-side `StopSequenceDecoder`, +/// which detokenizes worker output and trims the stop text. Existing +/// `stop_token_ids` are preserved and deduped. +fn encode_single_token_stops( + stops: Vec, + stop_token_ids: &mut Vec, + tokenizer: Option<&Arc>, +) { + // Without a tokenizer we cannot encode (not expected on paths that resolve + // one to tokenize the prompt). Safe: the strings are already dropped by the + // caller, so the router-side decoder remains the source of truth. + let Some(tokenizer) = tokenizer else { + if !stops.is_empty() { + warn!( + "No tokenizer available to encode string stop sequences; \ + relying on router-side stop decoder only" + ); + } + return; + }; + + for stop in stops { + if stop.is_empty() { + continue; + } + // add_special_tokens=false: we want the literal token(s) for the stop + // string, not a BOS/EOS-wrapped encoding. + match tokenizer.encode(&stop, false) { + Ok(encoding) => match encoding.token_ids() { + [id] => { + if !stop_token_ids.contains(id) { + stop_token_ids.push(*id); + } + } + ids => debug!( + stop = %stop, + token_count = ids.len(), + "string stop is not single-token; handled by router-side stop decoder" + ), + }, + Err(e) => warn!( + stop = %stop, + error = %e, + "Failed to encode string stop sequence; relying on router-side stop decoder" + ), + } + } +} + +/// Router-authoritative string-`stop` resolution for backends whose engine +/// cannot match string stops itself. +/// +/// vLLM over gRPC detokenizes server-side (`detokenize=bool(stop)`), TRT-LLM +/// tokenizes stop words server-side, and MLX has no string-`stop` field — those +/// keep their strings untouched. Two paths cannot: +/// - SGLang gRPC workers run with `skip_tokenizer_init=True` and reject string +/// stops outright (a 400 for any request carrying `stop`); and +/// - every direct-ZMQ backend (vLLM EngineCore, TokenSpeed) receives token ids +/// only, so the engine never sees — and cannot match — a stop string. +/// +/// For both, the router owns the tokenizer and already matches string stops via +/// `StopSequenceDecoder` (it detokenizes worker output and trims), so the worker +/// never needs the raw strings. This drops the string `stop` list and forwards +/// any single-token stop as a `stop_token_ids` entry for early stopping; the +/// router-side decoder handles the rest. This is the single resolution point +/// shared by SGLang gRPC and every ZMQ backend. +pub(crate) fn resolve_string_stops( + request: &mut ProtoGenerateRequest, + tokenizer: Option<&Arc>, + is_zmq: bool, +) { + // SGLang always needs it; the vLLM and TokenSpeed protos only when talking + // to a ZMQ backend (which is the sole path either reaches token-only). + match request { + ProtoGenerateRequest::Sglang(req) => { + if let Some(params) = req.sampling_params.as_mut() { + let stops = std::mem::take(&mut params.stop); + encode_single_token_stops(stops, &mut params.stop_token_ids, tokenizer); + } + } + ProtoGenerateRequest::Vllm(req) if is_zmq => { + if let Some(params) = req.sampling_params.as_mut() { + let stops = std::mem::take(&mut params.stop); + encode_single_token_stops(stops, &mut params.stop_token_ids, tokenizer); + // The engine only stops at EOS when the frontend supplies the + // ids, and the connect-time model-dir resolution has nothing + // to read when the worker's model id is a repo id rather than + // a local path. The tokenizer carries the merged EOS set, so + // fold it into the stop tokens as the always-available + // backstop — without it an uncapped request generates to the + // full context window. + if !params.ignore_eos { + if let Some(tokenizer) = tokenizer { + for &id in tokenizer.eos_token_ids() { + if !params.stop_token_ids.contains(&id) { + params.stop_token_ids.push(id); + } + } + } + } + } + } + ProtoGenerateRequest::TokenSpeed(req) if is_zmq => { + // TokenSpeed over ZMQ receives token ids only, so its wire + // translation drops raw `stop` strings; without this a single-token + // user stop would never reach the engine as a `stop_token_ids` + // entry. Resolve it exactly as the other token-only backends. No + // EOS fold here: unlike vLLM EngineCore, the TokenSpeed scheduler + // stops at EOS itself, so its translation carries no frontend ids. + if let Some(params) = req.sampling_params.as_mut() { + let stops = std::mem::take(&mut params.stop); + encode_single_token_stops(stops, &mut params.stop_token_ids, tokenizer); + } + } + _ => {} + } +} + /// Inject PD bootstrap metadata for SGLang if needed. /// /// SGLang uses DisaggregatedParams with bootstrap host/port/room. @@ -434,3 +559,222 @@ mod request_id_tests { assert!(id.starts_with("chatcmpl-")); } } + +#[cfg(test)] +mod stop_resolution_tests { + use std::sync::Arc; + + use llm_tokenizer::{mock::MockTokenizer, traits::Tokenizer}; + use smg_grpc_client::{sglang_proto, tokenspeed_proto, vllm_proto}; + + use super::{resolve_string_stops, ProtoGenerateRequest}; + + fn mock_tokenizer() -> Arc { + // MockTokenizer vocab: "." => 6, "Hello" => 1, "world" => 2. `encode` + // splits on whitespace, so "." => [6] (single) and "Hello world" => + // [1, 2] (multi); unknown words encode to []. + Arc::new(MockTokenizer::new()) + } + + fn sglang_request(stop: Vec<&str>, stop_token_ids: Vec) -> ProtoGenerateRequest { + ProtoGenerateRequest::Sglang(Box::new(sglang_proto::GenerateRequest { + sampling_params: Some(sglang_proto::SamplingParams { + stop: stop.into_iter().map(str::to_string).collect(), + stop_token_ids, + ..Default::default() + }), + ..Default::default() + })) + } + + fn vllm_request(stop: Vec<&str>, stop_token_ids: Vec) -> ProtoGenerateRequest { + ProtoGenerateRequest::Vllm(Box::new(vllm_proto::GenerateRequest { + sampling_params: Some(vllm_proto::SamplingParams { + stop: stop.into_iter().map(str::to_string).collect(), + stop_token_ids, + ..Default::default() + }), + ..Default::default() + })) + } + + fn tokenspeed_request(stop: Vec<&str>, stop_token_ids: Vec) -> ProtoGenerateRequest { + ProtoGenerateRequest::TokenSpeed(Box::new(tokenspeed_proto::GenerateRequest { + sampling_params: Some(tokenspeed_proto::SamplingParams { + stop: stop.into_iter().map(str::to_string).collect(), + stop_token_ids, + ..Default::default() + }), + ..Default::default() + })) + } + + fn tokenspeed_params(req: &ProtoGenerateRequest) -> &tokenspeed_proto::SamplingParams { + match req { + ProtoGenerateRequest::TokenSpeed(r) => r.sampling_params.as_ref().unwrap(), + _ => panic!("expected TokenSpeed request"), + } + } + + fn sglang_params(req: &ProtoGenerateRequest) -> &sglang_proto::SamplingParams { + match req { + ProtoGenerateRequest::Sglang(r) => r.sampling_params.as_ref().unwrap(), + _ => panic!("expected SGLang request"), + } + } + + fn vllm_params(req: &ProtoGenerateRequest) -> &vllm_proto::SamplingParams { + match req { + ProtoGenerateRequest::Vllm(r) => r.sampling_params.as_ref().unwrap(), + _ => panic!("expected vLLM request"), + } + } + + #[test] + fn sglang_single_token_becomes_stop_token_id() { + let mut req = sglang_request(vec!["."], vec![]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty(), "string stop should be cleared"); + assert_eq!(params.stop_token_ids, vec![6]); + } + + #[test] + fn sglang_multi_token_relies_on_router_decoder() { + // "Hello world" => [1, 2]: can't be a flat stop_token_id, so it must not + // be forwarded (would over-eagerly stop on any subtoken). + let mut req = sglang_request(vec!["Hello world"], vec![]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty()); + assert!(params.stop_token_ids.is_empty()); + } + + #[test] + fn sglang_mixed_only_single_token_forwarded_and_dedups() { + let mut req = sglang_request(vec![".", "Hello world"], vec![6, 42]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert!(params.stop.is_empty()); + assert_eq!( + params.stop_token_ids, + vec![6, 42], + "existing ids kept, no dup" + ); + } + + #[test] + fn sglang_without_tokenizer_still_clears_strings() { + let mut req = sglang_request(vec!["."], vec![]); + resolve_string_stops(&mut req, None, false); + + let params = sglang_params(&req); + assert!( + params.stop.is_empty(), + "strings dropped so worker won't 400" + ); + assert!(params.stop_token_ids.is_empty()); + } + + #[test] + fn vllm_resolved_only_over_zmq() { + // gRPC vLLM keeps its strings (the servicer detokenizes engine-side). + let mut grpc = vllm_request(vec!["."], vec![]); + resolve_string_stops(&mut grpc, Some(&mock_tokenizer()), false); + let params = vllm_params(&grpc); + assert_eq!( + params.stop, + vec![".".to_string()], + "gRPC vLLM stop preserved" + ); + assert!(params.stop_token_ids.is_empty()); + + // ZMQ vLLM (EngineCore sees token ids only) resolves like SGLang, and + // gains the tokenizer's EOS ids so generation always terminates. + let mut zmq = vllm_request(vec!["."], vec![]); + resolve_string_stops(&mut zmq, Some(&mock_tokenizer()), true); + let params = vllm_params(&zmq); + assert!(params.stop.is_empty(), "ZMQ vLLM stop cleared"); + assert_eq!(params.stop_token_ids, vec![6, 999]); + } + + #[test] + fn vllm_zmq_appends_tokenizer_eos_ids() { + let mut req = vllm_request(vec![], vec![7]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![7, 999]); + + // Already-present EOS ids are not duplicated. + let mut req = vllm_request(vec![], vec![999]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![999]); + } + + #[test] + fn vllm_zmq_ignore_eos_skips_injection() { + let mut req = ProtoGenerateRequest::Vllm(Box::new(vllm_proto::GenerateRequest { + sampling_params: Some(vllm_proto::SamplingParams { + stop_token_ids: vec![7], + ignore_eos: true, + ..Default::default() + }), + ..Default::default() + })); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!(vllm_params(&req).stop_token_ids, vec![7]); + } + + #[test] + fn noop_when_no_string_stops() { + let mut req = sglang_request(vec![], vec![7]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), false); + + let params = sglang_params(&req); + assert_eq!(params.stop_token_ids, vec![7], "unrelated ids untouched"); + } + + #[test] + fn tokenspeed_resolved_only_over_zmq() { + // A gRPC TokenSpeed request is never produced, but guard the gate: the + // strings must survive when is_zmq is false. + let mut grpc = tokenspeed_request(vec!["."], vec![]); + resolve_string_stops(&mut grpc, Some(&mock_tokenizer()), false); + let params = tokenspeed_params(&grpc); + assert_eq!(params.stop, vec![".".to_string()], "non-zmq stop preserved"); + assert!(params.stop_token_ids.is_empty()); + + // Over ZMQ the token-only wire drops raw strings, so a single-token stop + // must ride as a stop_token_ids entry instead. + let mut zmq = tokenspeed_request(vec!["."], vec![]); + resolve_string_stops(&mut zmq, Some(&mock_tokenizer()), true); + let params = tokenspeed_params(&zmq); + assert!(params.stop.is_empty(), "ZMQ TokenSpeed stop cleared"); + assert_eq!(params.stop_token_ids, vec![6]); + } + + #[test] + fn tokenspeed_zmq_does_not_fold_eos() { + // Unlike vLLM EngineCore, the TokenSpeed scheduler stops at EOS itself, + // so resolution must not append the tokenizer's EOS ids (999). + let mut req = tokenspeed_request(vec!["."], vec![]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + assert_eq!( + tokenspeed_params(&req).stop_token_ids, + vec![6], + "only the single-token stop, no EOS fold" + ); + } + + #[test] + fn tokenspeed_zmq_multi_token_relies_on_router_decoder() { + // "Hello world" => [1, 2]: not a flat stop id, so it must not forward. + let mut req = tokenspeed_request(vec!["Hello world"], vec![42]); + resolve_string_stops(&mut req, Some(&mock_tokenizer()), true); + let params = tokenspeed_params(&req); + assert!(params.stop.is_empty()); + assert_eq!(params.stop_token_ids, vec![42], "existing ids kept"); + } +} diff --git a/model_gateway/src/routers/grpc/regular/processor.rs b/model_gateway/src/routers/grpc/regular/processor.rs index 6d1d8d9e0..2275f3eb9 100644 --- a/model_gateway/src/routers/grpc/regular/processor.rs +++ b/model_gateway/src/routers/grpc/regular/processor.rs @@ -78,14 +78,19 @@ impl ResponseProcessor { // Accumulate text with early breaks let mut final_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => final_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { final_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -177,8 +182,14 @@ impl ResponseProcessor { } } - // Step 3: Use finish reason directly from proto (already OpenAI-compatible string) - let finish_reason_str = complete.finish_reason(); + // Step 3: Determine finish reason. A local stop-decoder match takes + // precedence over the engine's reason (which is "length" when stop + // strings are enforced gateway-side rather than by the backend). + let finish_reason_str = if stopped { + "stop" + } else { + complete.finish_reason() + }; // Override finish reason if we have tool calls let final_finish_reason_str = if tool_calls.is_some() { @@ -187,7 +198,12 @@ impl ResponseProcessor { finish_reason_str }; - let matched_stop = complete.matched_stop_json(); + // When the local decoder matched a stop string, surface it (the engine + // reports no stop_reason over the ZMQ path); otherwise use the engine's. + let matched_stop = stop_decoder + .matched_stop() + .map(|s| serde_json::Value::String(s.to_string())) + .or_else(|| complete.matched_stop_json()); // Step 4: Convert output logprobs if present let logprobs = complete.output_logprobs().map(|ref proto_logprobs| { @@ -579,14 +595,19 @@ impl ResponseProcessor { })?; let mut final_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => final_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { final_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -726,10 +747,20 @@ impl ResponseProcessor { } } - // Step 4: Determine stop_reason and stop_sequence (derived from same conditions) - let finish_reason_str = complete.finish_reason(); - let matched_stop = complete.matched_stop_json(); - let stop_sequence = matched_stop.and_then(|v| v.as_str().map(String::from)); + // Step 4: Determine stop_reason and stop_sequence (derived from same conditions). + // A local stop-decoder match takes precedence over the engine's reason + // (the backend has no stop-string detection over ZMQ), surfacing the + // matched sequence for a StopSequence result. + let finish_reason_str = if stopped { + "stop" + } else { + complete.finish_reason() + }; + let stop_sequence = stop_decoder.matched_stop().map(String::from).or_else(|| { + complete + .matched_stop_json() + .and_then(|v| v.as_str().map(String::from)) + }); let stop_reason = if tool_calls.is_some() || finish_reason_str == "tool_calls" { Some(messages::StopReason::ToolUse) @@ -832,14 +863,19 @@ impl ResponseProcessor { }; let mut decoded_text = String::new(); + let mut stopped = false; for output in outputs { match output { SequenceDecoderOutput::Text(t) => decoded_text.push_str(&t), SequenceDecoderOutput::StoppedWithText(t) => { decoded_text.push_str(&t); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; break; } - SequenceDecoderOutput::Stopped => break, SequenceDecoderOutput::Held => {} } } @@ -851,7 +887,12 @@ impl ResponseProcessor { prompt_tokens = prompt_tokens.max(complete.prompt_tokens()); total_completion += complete.completion_tokens(); - let finish_reason = { + // A local stop-decoder match takes precedence over the engine's + // reason (which is "length" when stop strings are enforced + // gateway-side rather than by the backend). + let finish_reason = if stopped { + Some("stop".to_string()) + } else { let reason = complete.finish_reason(); if reason.is_empty() { None @@ -868,7 +909,13 @@ impl ResponseProcessor { } }; - let matched_stop = complete.matched_stop_json(); + // When the local decoder matched a stop string, surface it (the + // engine reports no stop_reason over the ZMQ path); otherwise use + // the engine's. + let matched_stop = stop_decoder + .matched_stop() + .map(|s| serde_json::Value::String(s.to_string())) + .or_else(|| complete.matched_stop_json()); let suffix_len = completion_req.suffix.as_ref().map_or(0, |s| s.len()); let echo_len = if completion_req.echo { diff --git a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs index 4967044e9..3dd71e948 100644 --- a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs @@ -147,6 +147,13 @@ impl PipelineStage for ChatRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs index 6477ccecc..8bd9b63da 100644 --- a/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/completion/request_building.rs @@ -141,6 +141,12 @@ impl PipelineStage for CompletionRequestBuildingStage { let disaggregated = matches!(clients, ClientSelection::Disaggregated { .. }); let request_type = &ctx.input.request_type; let workers = ctx.state.workers.as_ref(); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + let tokenizer = ctx.tokenizer_arc(); let plan = match items.as_slice() { [] => { @@ -149,9 +155,8 @@ impl PipelineStage for CompletionRequestBuildingStage { "No prompts prepared", )) } - [item] => ExecutionPlan::generate( - self.plan_kind, - self.build_proto_request( + [item] => { + let mut proto_request = self.build_proto_request( builder_client, helpers::resolve_request_id( request_type, @@ -163,8 +168,10 @@ impl PipelineStage for CompletionRequestBuildingStage { &completion_request, request_type, workers, - )?, - ), + )?; + helpers::resolve_string_stops(&mut proto_request, tokenizer.as_ref(), is_zmq); + ExecutionPlan::generate(self.plan_kind, proto_request) + } batch_items => { // The shared id (client rid or middleware request id) stays // clean for the response; per-sub engine ids get a uniqueness @@ -186,14 +193,16 @@ impl PipelineStage for CompletionRequestBuildingStage { } else { format!("{shared_request_id}-p{i}") }; - requests.push(self.build_proto_request( + let mut proto_request = self.build_proto_request( builder_client, sub_request_id, item, &completion_request, request_type, workers, - )?); + )?; + helpers::resolve_string_stops(&mut proto_request, tokenizer.as_ref(), is_zmq); + requests.push(proto_request); } ExecutionPlan::Batch { kind: self.plan_kind, diff --git a/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs index 3be932e43..2f7634632 100644 --- a/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/generate/request_building.rs @@ -86,6 +86,13 @@ impl PipelineStage for GenerateRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs index 23d91b89a..aa73ff82a 100644 --- a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs @@ -148,6 +148,13 @@ impl PipelineStage for MessageRequestBuildingStage { ctx.state.workers.as_ref(), ); + // Resolve string `stop` sequences for engines that can't match them + // server-side (SGLang skip_tokenizer_init, and every direct-ZMQ + // backend): drop the strings, convert single-token stops to + // stop_token_ids; the router-side StopSequenceDecoder trims the text. + let is_zmq = builder_client.is_zmq(); + helpers::resolve_string_stops(&mut proto_request, ctx.tokenizer_arc().as_ref(), is_zmq); + if self.inject_pd_metadata { if let Some(workers) = ctx.state.workers.as_ref() { helpers::maybe_inject_pd_metadata(&mut proto_request, workers); diff --git a/model_gateway/src/routers/grpc/regular/streaming.rs b/model_gateway/src/routers/grpc/regular/streaming.rs index 505a8dc3b..b9d888f0c 100644 --- a/model_gateway/src/routers/grpc/regular/streaming.rs +++ b/model_gateway/src/routers/grpc/regular/streaming.rs @@ -271,6 +271,9 @@ impl StreamingProcessor { let mut stream_buffers: HashMap = HashMap::new(); let mut finish_reasons: HashMap = HashMap::new(); let mut matched_stops: HashMap> = HashMap::new(); + // Indices whose local stop decoder fired: their finish reason is pinned + // to "stop" and later engine output for the index is ignored. + let mut stopped_indices: HashSet = HashSet::new(); let mut prompt_tokens: HashMap = HashMap::new(); let mut completion_tokens = CompletionTokenTracker::new(); let mut cached_tokens: HashMap = HashMap::new(); @@ -384,6 +387,12 @@ impl StreamingProcessor { let index = chunk.index(); + // Once the local stop decoder has fired for an index, ignore + // any further engine output the backend emits for it. + if stopped_indices.contains(&index) { + continue; + } + completion_tokens.record_chunk(&chunk); // Get or create stop decoder for this index @@ -406,9 +415,26 @@ impl StreamingProcessor { }); // Process tokens through stop decoder - let (chunk_text, _should_stop) = + let (chunk_text, should_stop) = Self::process_chunk_tokens(stop_decoder, chunk.token_ids()); + if should_stop { + // Stop-decoder match takes precedence: pin "stop" even if + // the backend's eventual Complete carries "length" (the + // local stop sequence fired first). Any pre-stop text in + // `chunk_text` is still emitted below before the finish + // reason is flushed in Phase 4. + finish_reasons + .entry(index) + .or_insert_with(|| "stop".to_string()); + matched_stops.entry(index).or_insert_with(|| { + stop_decoder + .matched_stop() + .map(|s| Value::String(s.to_string())) + }); + stopped_indices.insert(index); + } + if chunk_text.is_empty() { continue; } @@ -568,9 +594,13 @@ impl StreamingProcessor { cached_tokens.insert(index, complete.cached_tokens()); reasoning_tokens.insert(index, complete.reasoning_tokens()); - finish_reasons.insert(index, complete.finish_reason().to_string()); - matched_stops.insert(index, complete.matched_stop_json()); + // A local stop-decoder match already pinned "stop" for this + // index; don't let the engine's finish reason overwrite it. + if !stopped_indices.contains(&index) { + finish_reasons.insert(index, complete.finish_reason().to_string()); + matched_stops.insert(index, complete.matched_stop_json()); + } // Don't break - continue reading all Complete messages for n>1 } @@ -1740,6 +1770,9 @@ impl StreamingProcessor { let mut prompt_tokens: u32 = 0; let mut finish_reason_str = String::new(); let mut matched_stop: Option = None; + // Set once the local stop decoder fires: pins "stop" and ignores later + // engine output (the backend has no stop-string detection over ZMQ). + let mut stopped = false; // Check parser availability once upfront. Run parser when the user explicitly // enabled thinking, or when the selected parser needs structural special tokens. @@ -1867,11 +1900,29 @@ impl StreamingProcessor { first_token_time = Some(Instant::now()); } + // Once the local stop decoder has fired, ignore further + // engine output for this (single-choice) request. + if stopped { + continue; + } + completion_tokens.record_chunk(&chunk); - let (chunk_text, _should_stop) = + let (chunk_text, should_stop) = Self::process_chunk_tokens(&mut stop_decoder, chunk.token_ids()); + if should_stop { + // Stop-decoder match takes precedence over the engine's + // eventual finish reason (the local stop sequence fired + // first). Pre-stop text in `chunk_text` is still emitted + // below; Phase 4 derives StopSequence from `matched_stop`. + stopped = true; + finish_reason_str = "stop".to_string(); + matched_stop = stop_decoder + .matched_stop() + .map(|s| Value::String(s.to_string())); + } + if chunk_text.is_empty() { continue; } @@ -2150,8 +2201,12 @@ impl StreamingProcessor { prompt_tokens = complete.prompt_tokens(); completion_tokens.record_complete(&complete); - finish_reason_str = complete.finish_reason().to_string(); - matched_stop = complete.matched_stop_json(); + // A local stop-decoder match already pinned "stop"; don't let + // the engine's finish reason overwrite it. + if !stopped { + finish_reason_str = complete.finish_reason().to_string(); + matched_stop = complete.matched_stop_json(); + } } ProtoResponseVariant::None => continue, } diff --git a/model_gateway/src/routers/grpc/zmq_client.rs b/model_gateway/src/routers/grpc/zmq_client.rs index d6165fa91..6e9890969 100644 --- a/model_gateway/src/routers/grpc/zmq_client.rs +++ b/model_gateway/src/routers/grpc/zmq_client.rs @@ -11,7 +11,8 @@ // the request-execution stage is reused unchanged. use std::{ - collections::HashMap, + collections::{BTreeSet, HashMap}, + path::Path, sync::Arc, time::{Duration, SystemTime, UNIX_EPOCH}, }; @@ -63,6 +64,69 @@ enum ZmqBackend { TokenSpeed(Arc), } +/// The model's EOS stop set, resolved from its local directory. EngineCore +/// has no tokenizer or model config — stopping at EOS is the frontend's job +/// (the ids ride each request), and without them generation only ends at +/// `max_tokens`. +#[derive(Debug, Clone, Default, PartialEq, Eq)] +pub struct EosTokenIds { + /// Primary EOS id, carried as the request's `_eos_token_id`. + primary: Option, + /// Extra EOS ids (multi-EOS models), merged into `stop_token_ids`. + extra: Vec, +} + +impl EosTokenIds { + pub fn new(primary: Option, extra: Vec) -> Self { + Self { primary, extra } + } + + /// Resolve from `config.json` + `generation_config.json` in a local model + /// directory: primary = the model config's first id, extras = every other + /// listed id. Missing files or fields degrade to fewer ids. + pub fn from_model_dir(dir: &Path) -> Self { + let model_ids = eos_ids_from_file(&dir.join("config.json")); + let gen_ids = eos_ids_from_file(&dir.join("generation_config.json")); + let primary = (model_ids.first().or_else(|| gen_ids.first())).copied(); + let mut extra = Vec::new(); + for id in model_ids.into_iter().chain(gen_ids) { + if Some(id) != primary && !extra.contains(&id) { + extra.push(id); + } + } + Self { primary, extra } + } +} + +/// Read a config file's `eos_token_id`, which is a single id or a list. +/// +/// A missing file is expected (a model ships `config.json`, +/// `generation_config.json`, or both), so read errors stay silent. A file +/// that exists but holds corrupt JSON is worth a `warn!`: it runs once at +/// connect time, and losing the EOS ids here silently manifests later as +/// generation running to `max_tokens`. +fn eos_ids_from_file(path: &Path) -> Vec { + let Ok(text) = std::fs::read_to_string(path) else { + return Vec::new(); + }; + match serde_json::from_str::(&text) { + Ok(config) => eos_ids_from_value(config.get("eos_token_id")), + Err(error) => { + tracing::warn!(path = %path.display(), %error, "failed to parse model config for EOS ids"); + Vec::new() + } + } +} + +fn eos_ids_from_value(value: Option<&serde_json::Value>) -> Vec { + let as_id = |v: &serde_json::Value| v.as_u64().and_then(|id| u32::try_from(id).ok()); + match value { + Some(serde_json::Value::Array(ids)) => ids.iter().filter_map(as_id).collect(), + Some(id) => as_id(id).into_iter().collect(), + None => Vec::new(), + } +} + /// Direct ZMQ connection to a same-host engine (vLLM EngineCore or TokenSpeed), /// presented behind the vLLM gRPC client surface. #[derive(Clone)] @@ -71,6 +135,9 @@ pub struct ZmqEngineClient { /// Model id advertised for metadata (the engine does not report it on the /// wire; it is configured at worker registration). model_id: String, + /// EOS ids attached to every vLLM request (the engine can't stop at EOS + /// without them). + eos: EosTokenIds, } impl ZmqEngineClient { @@ -81,12 +148,17 @@ impl ZmqEngineClient { /// engines connect to (chosen by SMG). `engine_count` is the number of DP /// ranks to await. `runtime` selects the wire protocol spoken over the shared /// transport (vLLM EngineCore vs TokenSpeed). + #[expect( + clippy::too_many_arguments, + reason = "transport constructor: endpoints, engine count, and runtime are all irreducible connection inputs" + )] pub async fn connect( handshake_address: &str, input_address: &str, output_address: &str, engine_count: usize, model_id: String, + eos: EosTokenIds, runtime: RuntimeType, timeout: Duration, ) -> Result> { @@ -133,7 +205,11 @@ impl ZmqEngineClient { // runtimes were rejected before the handshake. _ => ZmqBackend::Vllm(Arc::new(EngineCoreClient::new(transport))), }; - Ok(Self { backend, model_id }) + Ok(Self { + backend, + model_id, + eos, + }) } /// The engine runtime behind this connection (the wire protocol chosen at @@ -201,7 +277,7 @@ impl ZmqEngineClient { .and_then(|sp| sp.logprobs) .filter(|&n| n > 0) .map_or(0, |n| n as usize); - let request = translate_request(sub, max_model_len, model_dtype) + let request = translate_request(sub, max_model_len, model_dtype, &self.eos) .map_err(tonic::Status::invalid_argument)?; let stream = client.submit(request).await.map_err(zmq_status)?; streams.push(VllmGenerateStream::new(stream, index as u32, top_logprobs)); @@ -860,6 +936,7 @@ fn translate_request( req: vllm::GenerateRequest, max_model_len: u64, model_dtype: ModelDtype, + eos: &EosTokenIds, ) -> Result { let prompt_token_ids = match req.input { Some(vllm::generate_request::Input::Tokenized(tokenized)) => Some(tokenized.input_ids), @@ -907,7 +984,7 @@ fn translate_request( mm_features, sampling_params: req .sampling_params - .map(|sp| translate_sampling(sp, default_max_tokens)), + .map(|sp| translate_sampling(sp, default_max_tokens, eos)), arrival_time: now_secs(), data_parallel_rank, ..EngineCoreRequest::default() @@ -917,7 +994,23 @@ fn translate_request( fn translate_sampling( sp: vllm::SamplingParams, default_max_tokens: u32, + eos: &EosTokenIds, ) -> EngineCoreSamplingParams { + // Stopping at EOS is the frontend's duty here: the primary id rides + // `_eos_token_id`, extra ids merge into `stop_token_ids`, and the union + // feeds `_all_stop_token_ids` (engine-side `min_tokens` masking, built + // regardless of `ignore_eos`). + let mut stop_token_ids = sp.stop_token_ids; + if !sp.ignore_eos { + for id in &eos.extra { + if !stop_token_ids.contains(id) { + stop_token_ids.push(*id); + } + } + } + let mut all_stop_token_ids: BTreeSet = stop_token_ids.iter().copied().collect(); + all_stop_token_ids.extend(eos.primary); + all_stop_token_ids.extend(eos.extra.iter().copied()); let logit_bias = if sp.logit_bias.is_empty() { None } else { @@ -946,7 +1039,9 @@ fn translate_sampling( repetition_penalty: sp.repetition_penalty, max_tokens: sp.max_tokens.unwrap_or(default_max_tokens), min_tokens: sp.min_tokens, - stop_token_ids: sp.stop_token_ids, + stop_token_ids, + eos_token_id: (!sp.ignore_eos).then_some(eos.primary).flatten(), + all_stop_token_ids, seed: sp.seed.map(i64::from), logprobs: sp.logprobs, // prompt_logprobs is rejected in `translate_request` (no renderer @@ -1032,8 +1127,6 @@ fn zmq_status(error: engine_zmq_client::Error) -> tonic::Status { #[cfg(test)] mod tests { - use std::collections::BTreeSet; - use engine_zmq_client::{ mock_engine::{connect_to_frontend, default_ready_response, EngineInbound}, protocol::vllm::{ @@ -1092,6 +1185,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1213,6 +1307,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1360,6 +1455,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::TokenSpeed, Duration::from_secs(10) ), @@ -1682,6 +1778,7 @@ mod tests { }), 4096, ModelDtype::BFloat16, + &EosTokenIds::default(), ) .expect_err("prompt logprobs rejected"); assert!(err.contains("prompt logprobs"), "{err}"); @@ -1690,11 +1787,16 @@ mod tests { #[test] fn vllm_defaults_unset_max_tokens_to_remaining_context() { let max_tokens = |sampling, max_model_len| { - translate_request(tokenized_req(sampling), max_model_len, ModelDtype::BFloat16) - .expect("request translated") - .sampling_params - .expect("sampling params present") - .max_tokens + translate_request( + tokenized_req(sampling), + max_model_len, + ModelDtype::BFloat16, + &EosTokenIds::default(), + ) + .expect("request translated") + .sampling_params + .expect("sampling params present") + .max_tokens }; // Unset max_tokens defaults to `max_model_len - prompt_len` (prompt is @@ -1713,6 +1815,60 @@ mod tests { ); } + #[test] + fn vllm_attaches_eos_stop_ids() { + let eos = EosTokenIds::new(Some(5), vec![7]); + let sampling = |sp| { + translate_request(tokenized_req(sp), 4096, ModelDtype::BFloat16, &eos) + .expect("request translated") + .sampling_params + .expect("sampling params present") + }; + + // Primary rides `_eos_token_id`, extras merge into `stop_token_ids` + // without duplicating, and the union lands in `_all_stop_token_ids`. + let sp = sampling(vllm::SamplingParams { + stop_token_ids: vec![7, 9], + ..Default::default() + }); + assert_eq!(sp.eos_token_id, Some(5)); + assert_eq!(sp.stop_token_ids, vec![7, 9]); + assert_eq!(sp.all_stop_token_ids, BTreeSet::from([5, 7, 9])); + + // ignore_eos drops the EOS stops from the wire but keeps the + // bookkeeping set (mirrors the reference frontend). + let sp = sampling(vllm::SamplingParams { + stop_token_ids: vec![9], + ignore_eos: true, + ..Default::default() + }); + assert_eq!(sp.eos_token_id, None); + assert_eq!(sp.stop_token_ids, vec![9]); + assert_eq!(sp.all_stop_token_ids, BTreeSet::from([5, 7, 9])); + } + + #[test] + fn eos_token_ids_resolve_from_model_dir() { + let dir = tempfile::tempdir().expect("tempdir"); + std::fs::write(dir.path().join("config.json"), r#"{"eos_token_id": 5}"#).unwrap(); + std::fs::write( + dir.path().join("generation_config.json"), + r#"{"eos_token_id": [5, 7, 9]}"#, + ) + .unwrap(); + assert_eq!( + EosTokenIds::from_model_dir(dir.path()), + EosTokenIds::new(Some(5), vec![7, 9]), + ); + + // Missing files degrade to no ids, not an error. + let empty = tempfile::tempdir().expect("tempdir"); + assert_eq!( + EosTokenIds::from_model_dir(empty.path()), + EosTokenIds::default(), + ); + } + #[test] fn vllm_translates_structured_output_constraints() { use engine_zmq_client::protocol::vllm::structured_outputs::{ @@ -1727,6 +1883,7 @@ mod tests { }), 4096, ModelDtype::BFloat16, + &EosTokenIds::default(), ) .expect("constraint translated") .sampling_params @@ -1840,6 +1997,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), @@ -1940,6 +2098,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::TokenSpeed, Duration::from_secs(10) ), @@ -2043,6 +2202,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ), diff --git a/model_gateway/src/worker/worker.rs b/model_gateway/src/worker/worker.rs index a117fc545..d34b9a4e3 100644 --- a/model_gateway/src/worker/worker.rs +++ b/model_gateway/src/worker/worker.rs @@ -35,7 +35,7 @@ use crate::{ grpc::{ backend_client::BackendClient, client::GrpcClient, - zmq_client::{ZmqEngineClient, ZMQ_LOOPBACK_HOST}, + zmq_client::{EosTokenIds, ZmqEngineClient, ZMQ_LOOPBACK_HOST}, }, }, }; @@ -200,6 +200,19 @@ async fn connect_zmq_backend( let (handshake, input, output) = zmq_socket_addresses(&base_url, handshake_override.as_deref())?; ensure_ipc_socket_dir(&base_url).await?; + // The engine can't stop at EOS on its own (it has no tokenizer or model + // config); resolve the EOS ids from the local model dir so every request + // carries them. + let model_dir = std::path::Path::new(&model_id); + let eos = if model_dir.is_dir() { + EosTokenIds::from_model_dir(model_dir) + } else { + tracing::warn!( + "ZMQ worker model id '{model_id}' is not a local model directory; EOS ids \ + unavailable — generation stops only at max_tokens or explicit stops" + ); + EosTokenIds::default() + }; tracing::info!("Binding ZMQ client for worker {base_url} (handshake={handshake})"); match ZmqEngineClient::connect( &handshake, @@ -207,6 +220,7 @@ async fn connect_zmq_backend( &output, 1, model_id, + eos, runtime, ZMQ_CONNECT_TIMEOUT, ) @@ -2826,6 +2840,7 @@ mod tests { &output, 1, "m".to_string(), + EosTokenIds::default(), RuntimeType::Vllm, Duration::from_secs(10) ),