[Rust Frontend] add repetition_detection support to sampling params (#46684)

Co-authored-by: Bugen Zhao <i@bugenzhao.com>
Signed-off-by: Bugen Zhao <i@bugenzhao.com>

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`:
+/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L109-L144>
+#[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<RepetitionDetectionParams>,
     /// Token IDs that stop generation.
     pub stop_token_ids: Vec<u32>,
     /// 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);

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(),

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<StopReason>),
 }
 
 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<StopReason> {
         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),
     })
 }
 

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)
         }
     };

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");
         }

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)?,

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<HashMap<String, Value>>,
 
     /// Parameters for detecting repetitive N-gram patterns in output tokens
-    pub repetition_detection: Option<Value>,
+    pub repetition_detection: Option<RepetitionDetectionParams>,
 }
 
 impl Default for ChatCompletionRequest {

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(())
 }

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<CompletionStreamResponse, ApiError> {
-    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!(

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)?,

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<f32>,
 
+    /// Parameters for detecting repetitive N-gram patterns in output tokens
+    pub repetition_detection: Option<RepetitionDetectionParams>,
+
     /// Length penalty for beam search
     pub length_penalty: Option<f32>,
 

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,

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<i64>) -> Result<Option<u64>> {
     }
 }
 
+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: {},

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<f32>,
+    /// Parameters for detecting repetitive N-gram patterns. `None` means no
+    /// explicit user override.
+    pub repetition_detection: Option<RepetitionDetectionParams>,
     /// Explicit stop token IDs provided by the caller. `None` means no explicit
     /// user override.
     pub stop_token_ids: Option<Vec<u32>>,
@@ -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,