From 6429d5f527f57d4a88b79b7fda2fcdd32f4a0e13 Mon Sep 17 00:00:00 2001 From: yangyang-cs95 <360635545@qq.com> Date: Fri, 3 Jul 2026 14:06:23 +0800 Subject: [PATCH] [Rust Frontend] add repetition_detection support to sampling params (#46684) Co-authored-by: Bugen Zhao Signed-off-by: Bugen Zhao --- .../src/protocol/sampling.rs | 29 ++++++ .../engine-core-client/src/tests/client.rs | 1 + rust/src/llm/src/output.rs | 8 +- rust/src/server/src/grpc/convert.rs | 2 +- .../src/routes/openai/chat_completions.rs | 2 +- .../routes/openai/chat_completions/convert.rs | 1 + .../routes/openai/chat_completions/types.rs | 3 +- .../openai/chat_completions/validate.rs | 5 - .../server/src/routes/openai/completions.rs | 40 ++++++-- .../src/routes/openai/completions/convert.rs | 1 + .../src/routes/openai/completions/types.rs | 4 + rust/src/text/src/error.rs | 3 + rust/src/text/src/lower.rs | 93 ++++++++++++++++++- rust/src/text/src/request.rs | 5 + 14 files changed, 177 insertions(+), 20 deletions(-) diff --git a/rust/src/engine-core-client/src/protocol/sampling.rs b/rust/src/engine-core-client/src/protocol/sampling.rs index b724b36dc5..5b52ded065 100644 --- a/rust/src/engine-core-client/src/protocol/sampling.rs +++ b/rust/src/engine-core-client/src/protocol/sampling.rs @@ -21,6 +21,30 @@ fn default_max_tokens() -> u32 { 16 } +/// +/// Parameters for detecting repetitive N-gram patterns in output tokens. +/// +/// Mirrors Python's `RepetitionDetectionParams`: +/// +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct RepetitionDetectionParams { + /// Maximum N-gram size to check. 0 disables detection. + pub max_pattern_size: u32, + /// Minimum N-gram size to check. Defaults to 1 when zero. + #[serde(default)] + pub min_pattern_size: u32, + /// Minimum number of repetitions to trigger detection (must be >= 2). + pub min_count: u32, +} + +impl RepetitionDetectionParams { + /// Return `true` when the params are effectively disabled (max_pattern_size + /// is 0). + pub fn is_disabled(&self) -> bool { + self.max_pattern_size == 0 + } +} + /// Engine-core-facing sampling parameters for text generation. /// /// This is the normalized southbound subset used by the Rust frontend when it @@ -76,6 +100,9 @@ pub struct EngineCoreSamplingParams { /// Repetition penalty applied by the sampler. #[serde(default = "default_repetition_penalty")] pub repetition_penalty: f32, + /// Parameters for detecting repetitive N-gram patterns. `None` disables + /// detection. + pub repetition_detection: Option, /// Token IDs that stop generation. pub stop_token_ids: Vec, /// Primary EOS token ID used by engine-core's dedicated EOS stop path. @@ -134,6 +161,7 @@ impl EngineCoreSamplingParams { frequency_penalty: 0.0, presence_penalty: 0.0, repetition_penalty: 1.0, + repetition_detection: None, stop_token_ids: Vec::new(), eos_token_id: None, all_stop_token_ids: BTreeSet::new(), @@ -203,6 +231,7 @@ mod tests { assert_eq!(sampling.frequency_penalty, 0.0); assert_eq!(sampling.presence_penalty, 0.0); assert_eq!(sampling.repetition_penalty, 1.0); + assert_eq!(sampling.repetition_detection, None); assert_eq!(sampling.logprobs, None); assert_eq!(sampling.prompt_logprobs, None); assert_eq!(sampling.eos_token_id, None); diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 00aae6fe0c..b433e90603 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -2462,6 +2462,7 @@ fn python_msgpack_fixtures_match_rust_encoding() { frequency_penalty: 0.0, presence_penalty: 0.0, repetition_penalty: 1.0, + repetition_detection: None, stop_token_ids: Vec::new(), eos_token_id: None, all_stop_token_ids: BTreeSet::new(), diff --git a/rust/src/llm/src/output.rs b/rust/src/llm/src/output.rs index 7602df5b80..609088ec6b 100644 --- a/rust/src/llm/src/output.rs +++ b/rust/src/llm/src/output.rs @@ -69,7 +69,7 @@ pub enum FinishReason { /// A retryable request-level internal error occurred. Error, /// A repetitive token pattern was detected. - Repetition, + Repetition(Option), } impl FinishReason { @@ -87,7 +87,7 @@ impl FinishReason { Self::Length => "length", Self::Abort => "abort", Self::Error => "error", - Self::Repetition => "repetition", + Self::Repetition(_) => "repetition", } } @@ -96,6 +96,7 @@ impl FinishReason { pub fn as_stop_reason(&self) -> Option<&StopReason> { match self { Self::Stop(stop_reason) => stop_reason.as_ref(), + Self::Repetition(stop_reason) => stop_reason.as_ref(), _ => None, } } @@ -105,6 +106,7 @@ impl FinishReason { pub fn into_stop_reason(self) -> Option { match self { Self::Stop(stop_reason) => stop_reason, + Self::Repetition(stop_reason) => stop_reason, _ => None, } } @@ -119,7 +121,7 @@ fn finish_reason_from_engine( EngineCoreFinishReason::Length => FinishReason::Length, EngineCoreFinishReason::Abort => FinishReason::Abort, EngineCoreFinishReason::Error => FinishReason::Error, - EngineCoreFinishReason::Repetition => FinishReason::Repetition, + EngineCoreFinishReason::Repetition => FinishReason::Repetition(stop_reason), }) } diff --git a/rust/src/server/src/grpc/convert.rs b/rust/src/server/src/grpc/convert.rs index 5dfeba08a5..9de8972720 100644 --- a/rust/src/server/src/grpc/convert.rs +++ b/rust/src/server/src/grpc/convert.rs @@ -332,7 +332,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo { (PbFinishReason::Stop as i32, sr) } FinishReason::Length => (PbFinishReason::Length as i32, None), - FinishReason::Abort | FinishReason::Error | FinishReason::Repetition => { + FinishReason::Abort | FinishReason::Error | FinishReason::Repetition(_) => { (PbFinishReason::Aborted as i32, None) } }; diff --git a/rust/src/server/src/routes/openai/chat_completions.rs b/rust/src/server/src/routes/openai/chat_completions.rs index 5a879267c9..0429d9ac8b 100644 --- a/rust/src/server/src/routes/openai/chat_completions.rs +++ b/rust/src/server/src/routes/openai/chat_completions.rs @@ -805,7 +805,7 @@ fn chat_finish_reason_to_openai( FinishReason::Stop(_) => Ok("stop"), FinishReason::Length => Ok("length"), FinishReason::Abort => Ok("abort"), - FinishReason::Repetition => Ok("stop"), + FinishReason::Repetition(_) => Ok("repetition"), FinishReason::Error => { bail_server_error!("Internal server error"); } diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 2d64d194bf..26a7627149 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -122,6 +122,7 @@ pub(super) fn prepare_chat_request( frequency_penalty: request.frequency_penalty, presence_penalty: request.presence_penalty, repetition_penalty: request.repetition_penalty, + repetition_detection: request.repetition_detection, stop_token_ids: request.stop_token_ids, ignore_eos: request.ignore_eos, logit_bias: convert_logit_bias(request.logit_bias)?, diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index a76a52b8ad..b5681d0abd 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -6,6 +6,7 @@ use serde_json::Value; use serde_with::SerializeDisplay; use validator::Validate; use vllm_chat::ReasoningEffort; +use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ @@ -236,7 +237,7 @@ pub struct ChatCompletionRequest { pub vllm_xargs: Option>, /// Parameters for detecting repetitive N-gram patterns in output tokens - pub repetition_detection: Option, + pub repetition_detection: Option, } impl Default for ChatCompletionRequest { diff --git a/rust/src/server/src/routes/openai/chat_completions/validate.rs b/rust/src/server/src/routes/openai/chat_completions/validate.rs index f731cb2cde..2859b02274 100644 --- a/rust/src/server/src/routes/openai/chat_completions/validate.rs +++ b/rust/src/server/src/routes/openai/chat_completions/validate.rs @@ -94,11 +94,6 @@ pub(super) fn validate_request_compat( "mm_processor_kwargs", "mm_processor_kwargs is not supported.", )?; - reject_non_default( - request.repetition_detection.as_ref(), - "repetition_detection", - "repetition_detection is not supported.", - )?; Ok(()) } diff --git a/rust/src/server/src/routes/openai/completions.rs b/rust/src/server/src/routes/openai/completions.rs index 462e575aeb..0dd0eadf33 100644 --- a/rust/src/server/src/routes/openai/completions.rs +++ b/rust/src/server/src/routes/openai/completions.rs @@ -14,9 +14,11 @@ use axum::http::HeaderMap; use axum::response::sse::{Event, Sse}; use axum::response::{IntoResponse, Response}; use futures::{Stream, StreamExt as _, pin_mut}; +use serde_json::Value; use thiserror_ext::AsReport as _; use tracing::{debug, error, info, trace}; use tracing_futures::Instrument as _; +use vllm_engine_core_client::protocol::output::StopReason; use vllm_text::{ DecodedPromptLogprobs, DecodedTextEvent, FinishReason, TextOutputStream, TextOutputStreamExt as _, @@ -142,9 +144,7 @@ async fn collect_completion( .await .map_err(|error| server_error!("completion stream failed: {}", error.to_report_string()))?; let finish_reason = collected.finish_reason.clone(); - let stop_reason = finish_reason - .as_stop_reason() - .map(|sr| serde_json::to_value(sr).expect("StopReason must serialize to JSON")); + let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json); let prompt_char_count = echo.as_ref().map(|prompt| text_len(prompt)).unwrap_or_default(); let logprobs = if requested_logprobs.is_some() && prompt_only { @@ -181,7 +181,7 @@ async fn collect_completion( Some(prompt) if prompt_only => prompt.clone(), Some(prompt) => format!("{prompt}{}", collected.text), }; - let finish_reason = completion_finish_reason_to_openai(finish_reason)?.to_string(); + let finish_reason = completion_finish_reason_to_openai(&finish_reason)?.to_string(); let usage = Usage::from_token_usage(collected.usage, enable_prompt_tokens_details); if enable_log_requests { @@ -441,21 +441,24 @@ fn final_chunk( created: u64, finish_reason: FinishReason, ) -> Result { - let finish_reason = completion_finish_reason_to_openai(finish_reason)?; + let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json); + let finish_reason = completion_finish_reason_to_openai(&finish_reason)?; let mut chunk = CompletionStreamResponse::new(request_id, response_model, created); chunk.choices.push(CompletionStreamChoice { finish_reason: Some(finish_reason.to_string()), + stop_reason, ..Default::default() }); Ok(chunk) } fn completion_finish_reason_to_openai( - finish_reason: FinishReason, + finish_reason: &FinishReason, ) -> Result<&'static str, ApiError> { match finish_reason { - FinishReason::Stop(_) | FinishReason::Repetition => Ok("stop"), + FinishReason::Stop(_) => Ok("stop"), + FinishReason::Repetition(_) => Ok("repetition"), FinishReason::Length => Ok("length"), FinishReason::Abort => Ok("abort"), FinishReason::Error => { @@ -464,6 +467,10 @@ fn completion_finish_reason_to_openai( } } +fn stop_reason_to_json(stop_reason: &StopReason) -> Value { + serde_json::to_value(stop_reason).expect("StopReason must serialize to JSON") +} + fn prompt_only_logprobs_to_openai( prompt_logprobs: Option<&DecodedPromptLogprobs>, prompt: &str, @@ -577,6 +584,7 @@ fn done_sse_event() -> Event { mod tests { use futures::{StreamExt as _, stream}; use itertools::Itertools as _; + use vllm_engine_core_client::protocol::output::StopReason; use vllm_text::{ DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTextEvent, DecodedTokenLogprob, FinishReason, Finished, @@ -670,7 +678,9 @@ mod tests { output_token_count: 2, cached_token_count: 3, }, - finish_reason: FinishReason::stop_eos(), + finish_reason: FinishReason::Repetition(Some(StopReason::Text( + "repetition_detected".to_string(), + ))), kv_transfer_params: None, }), }), @@ -726,6 +736,20 @@ mod tests { CompletionSseChunk::Usage(_) => panic!("expected regular chunk"), } + match &chunks[2] { + CompletionSseChunk::Chunk(chunk) => { + assert_eq!( + chunk.choices[0].finish_reason.as_deref(), + Some("repetition") + ); + assert_eq!( + chunk.choices[0].stop_reason, + Some(serde_json::json!("repetition_detected")) + ); + } + CompletionSseChunk::Usage(_) => panic!("expected regular chunk"), + } + match &chunks[3] { CompletionSseChunk::Usage(chunk) => { assert_eq!( diff --git a/rust/src/server/src/routes/openai/completions/convert.rs b/rust/src/server/src/routes/openai/completions/convert.rs index 8ea86c0f48..fa80e16e0f 100644 --- a/rust/src/server/src/routes/openai/completions/convert.rs +++ b/rust/src/server/src/routes/openai/completions/convert.rs @@ -117,6 +117,7 @@ pub(super) fn prepare_completion_request( frequency_penalty: request.frequency_penalty, presence_penalty: request.presence_penalty, repetition_penalty: request.repetition_penalty, + repetition_detection: request.repetition_detection, stop_token_ids: request.stop_token_ids, ignore_eos: request.ignore_eos, logit_bias: convert_logit_bias(request.logit_bias)?, diff --git a/rust/src/server/src/routes/openai/completions/types.rs b/rust/src/server/src/routes/openai/completions/types.rs index 153730ab1b..fe637d82de 100644 --- a/rust/src/server/src/routes/openai/completions/types.rs +++ b/rust/src/server/src/routes/openai/completions/types.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; +use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use vllm_text::Prompt; use crate::routes::openai::utils::types::{ @@ -100,6 +101,9 @@ pub struct CompletionRequest { /// Repetition penalty for reducing repetitive text pub repetition_penalty: Option, + /// Parameters for detecting repetitive N-gram patterns in output tokens + pub repetition_detection: Option, + /// Length penalty for beam search pub length_penalty: Option, diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index 96ddced584..76827ae581 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -27,6 +27,8 @@ pub enum Error { MinTokensExceedsMaxTokens { min_tokens: u32, max_tokens: u32 }, #[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")] InvalidThinkingTokenBudget, + #[error("invalid repetition detection params: {message}")] + InvalidRepetitionDetection { message: String }, #[error("text request stream `{request_id}` closed before terminal output")] StreamClosedBeforeTerminalOutput { request_id: String }, #[error(transparent)] @@ -47,6 +49,7 @@ impl Error { | Self::TokenIds(_) | Self::MinTokensExceedsMaxTokens { .. } | Self::InvalidThinkingTokenBudget + | Self::InvalidRepetitionDetection { .. } // An empty tokenized prompt detected later, at request prepare // time, surfaces through the transparent Llm wrapper. | Self::Llm(LlmError::EmptyPromptTokenIds { .. }) => true, diff --git a/rust/src/text/src/lower.rs b/rust/src/text/src/lower.rs index 31cf0058cd..eabece37f0 100644 --- a/rust/src/text/src/lower.rs +++ b/rust/src/text/src/lower.rs @@ -5,7 +5,9 @@ pub(crate) mod token_ids; use logprobs::validate_logprobs; use token_ids::{validate_prompt_token_ids, validate_vocab_range}; -use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams; +use vllm_engine_core_client::protocol::sampling::{ + EngineCoreSamplingParams, RepetitionDetectionParams, +}; use vllm_llm::GenerateRequest; use vllm_tokenizer::Tokenizer; @@ -93,6 +95,7 @@ pub fn lower_sampling_params( frequency_penalty, presence_penalty, repetition_penalty, + repetition_detection, stop_token_ids, ignore_eos, logit_bias, @@ -110,6 +113,7 @@ pub fn lower_sampling_params( logprob_token_ids.as_deref(), sampling_limits, )?; + validate_repetition_detection(repetition_detection.as_ref())?; // Mirrors the model-generation-config inheritance used by vLLM's OpenAI chat // path: https://github.com/vllm-project/vllm/blob/bc2c0c86efb28e77677a3cfb8687e976914a313a/vllm/entrypoints/openai/chat_completion/protocol.py#L424-L450 @@ -163,6 +167,7 @@ pub fn lower_sampling_params( frequency_penalty, presence_penalty, repetition_penalty, + repetition_detection: repetition_detection.filter(|p| !p.is_disabled()), stop_token_ids, eos_token_id: (!ignore_eos).then_some(primary_eos_token_id).flatten(), all_stop_token_ids, @@ -194,6 +199,33 @@ fn normalize_thinking_token_budget(value: Option) -> Result> { } } +fn validate_repetition_detection(params: Option<&RepetitionDetectionParams>) -> Result<()> { + let Some(params) = params else { + return Ok(()); + }; + + if params.min_pattern_size > params.max_pattern_size { + return Err(Error::InvalidRepetitionDetection { + message: format!( + "`min_pattern_size` must be less than or equal to \ + `max_pattern_size`, got min_pattern_size={}, \ + max_pattern_size={}", + params.min_pattern_size, params.max_pattern_size + ), + }); + } + if params.max_pattern_size > 0 && params.min_count < 2 { + return Err(Error::InvalidRepetitionDetection { + message: format!( + "`min_count` must be at least 2, got min_count={}", + params.min_count + ), + }); + } + + Ok(()) +} + /// Convert bad-word strings into token-ID sequences, following the Python vLLM /// logic in `SamplingParams.update_from_tokenizer()`. /// @@ -390,6 +422,59 @@ mod tests { )); } + #[test] + fn lower_sampling_params_validates_repetition_detection() { + let lower = |repetition_detection| { + lower_sampling_params_with_limits( + SamplingParams { + repetition_detection, + ..SamplingParams::default() + }, + sample_sampling_limits(), + ) + }; + + let enabled = RepetitionDetectionParams { + max_pattern_size: 4, + min_pattern_size: 2, + min_count: 2, + }; + assert_eq!( + lower(Some(enabled.clone())).unwrap().repetition_detection, + Some(enabled) + ); + + let disabled = RepetitionDetectionParams { + max_pattern_size: 0, + min_pattern_size: 0, + min_count: 0, + }; + assert_eq!(lower(Some(disabled)).unwrap().repetition_detection, None); + + let error = lower(Some(RepetitionDetectionParams { + max_pattern_size: 1, + min_pattern_size: 2, + min_count: 2, + })) + .unwrap_err(); + let Error::InvalidRepetitionDetection { message } = error else { + panic!("expected repetition_detection validation error"); + }; + assert!(message.contains("min_pattern_size=2")); + assert!(message.contains("max_pattern_size=1")); + + let error = lower(Some(RepetitionDetectionParams { + max_pattern_size: 1, + min_pattern_size: 1, + min_count: 1, + })) + .unwrap_err(); + let Error::InvalidRepetitionDetection { message } = error else { + panic!("expected repetition_detection validation error"); + }; + assert!(message.contains("min_count=1")); + } + #[test] fn lower_text_request_applies_python_style_eos_hints() { let prepared = lower_text_request( @@ -417,6 +502,7 @@ mod tests { frequency_penalty: 0.0, presence_penalty: 0.0, repetition_penalty: 1.0, + repetition_detection: None, stop_token_ids: [ 77, ], @@ -469,6 +555,7 @@ mod tests { frequency_penalty: 0.0, presence_penalty: 0.0, repetition_penalty: 1.0, + repetition_detection: None, stop_token_ids: [], eos_token_id: None, all_stop_token_ids: { @@ -600,6 +687,7 @@ mod tests { frequency_penalty: 0.0, presence_penalty: 0.0, repetition_penalty: 1.0, + repetition_detection: None, stop_token_ids: [ 151643, ], @@ -662,6 +750,7 @@ mod tests { frequency_penalty: 0.0, presence_penalty: 0.0, repetition_penalty: 1.0, + repetition_detection: None, stop_token_ids: [ 11, 77, @@ -732,6 +821,7 @@ mod tests { frequency_penalty: 0.0, presence_penalty: 0.0, repetition_penalty: 1.2, + repetition_detection: None, stop_token_ids: [], eos_token_id: None, all_stop_token_ids: {}, @@ -980,6 +1070,7 @@ mod tests { frequency_penalty: 0.0, presence_penalty: 0.0, repetition_penalty: 1.2, + repetition_detection: None, stop_token_ids: [], eos_token_id: None, all_stop_token_ids: {}, diff --git a/rust/src/text/src/request.rs b/rust/src/text/src/request.rs index 3b59b687b5..c64e9ca05e 100644 --- a/rust/src/text/src/request.rs +++ b/rust/src/text/src/request.rs @@ -6,6 +6,7 @@ use serde_json::Value; use vllm_engine_core_client::protocol::lora::LoraRequest; use vllm_engine_core_client::protocol::multimodal::MmFeatures; use vllm_engine_core_client::protocol::request::ReasoningParserKwargs; +use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams; use crate::error::{Error, Result}; @@ -83,6 +84,9 @@ pub struct SamplingParams { /// Repetition penalty applied by the sampler. `None` means no explicit user /// override. pub repetition_penalty: Option, + /// Parameters for detecting repetitive N-gram patterns. `None` means no + /// explicit user override. + pub repetition_detection: Option, /// Explicit stop token IDs provided by the caller. `None` means no explicit /// user override. pub stop_token_ids: Option>, @@ -130,6 +134,7 @@ impl Default for SamplingParams { frequency_penalty: None, presence_penalty: None, repetition_penalty: None, + repetition_detection: None, stop_token_ids: None, ignore_eos: false, logit_bias: None,