@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 6429d5f527

[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>

yangyang-cs95 committed Jul 3, 2026 at 14:06 UTC 6429d5f527f57d4a88b79b7fda2fcdd32f4a0e13

14 files changed +177 -20

rust/src/engine-core-client/src/protocol/sampling.rs

+29

@@ -21,6 +21,30 @@ fn default_max_tokens() -> u32 {

21 16

22 }

24 +///

25 +/// Parameters for detecting repetitive N-gram patterns in output tokens.

26 +///

27 +/// Mirrors Python's `RepetitionDetectionParams`:

28 +/// <https://github.com/vllm-project/vllm/blob/f22d6e026798a74e6542a52ef776c054f2de572a/vllm/sampling_params.py#L109-L144>

29 +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]

30 +pub struct RepetitionDetectionParams {

31 + /// Maximum N-gram size to check. 0 disables detection.

32 + pub max_pattern_size: u32,

33 + /// Minimum N-gram size to check. Defaults to 1 when zero.

34 + #[serde(default)]

35 + pub min_pattern_size: u32,

36 + /// Minimum number of repetitions to trigger detection (must be >= 2).

37 + pub min_count: u32,

38 +}

39 +

40 +impl RepetitionDetectionParams {

41 + /// Return `true` when the params are effectively disabled (max_pattern_size

42 + /// is 0).

43 + pub fn is_disabled(&self) -> bool {

44 + self.max_pattern_size == 0

45 + }

46 +}

47 +

48 /// Engine-core-facing sampling parameters for text generation.

49 ///

50 /// This is the normalized southbound subset used by the Rust frontend when it

@@ -76,6 +100,9 @@ pub struct EngineCoreSamplingParams {

100 /// Repetition penalty applied by the sampler.

101 #[serde(default = "default_repetition_penalty")]

102 pub repetition_penalty: f32,

103 + /// Parameters for detecting repetitive N-gram patterns. `None` disables

104 + /// detection.

105 + pub repetition_detection: Option<RepetitionDetectionParams>,

106 /// Token IDs that stop generation.

107 pub stop_token_ids: Vec<u32>,

108 /// Primary EOS token ID used by engine-core's dedicated EOS stop path.

@@ -134,6 +161,7 @@ impl EngineCoreSamplingParams {

161 frequency_penalty: 0.0,

162 presence_penalty: 0.0,

163 repetition_penalty: 1.0,

164 + repetition_detection: None,

165 stop_token_ids: Vec::new(),

166 eos_token_id: None,

167 all_stop_token_ids: BTreeSet::new(),

@@ -203,6 +231,7 @@ mod tests {

231 assert_eq!(sampling.frequency_penalty, 0.0);

232 assert_eq!(sampling.presence_penalty, 0.0);

233 assert_eq!(sampling.repetition_penalty, 1.0);

234 + assert_eq!(sampling.repetition_detection, None);

235 assert_eq!(sampling.logprobs, None);

236 assert_eq!(sampling.prompt_logprobs, None);

237 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() {

2462 frequency_penalty: 0.0,

2463 presence_penalty: 0.0,

2464 repetition_penalty: 1.0,

2465 + repetition_detection: None,

2466 stop_token_ids: Vec::new(),

2467 eos_token_id: None,

2468 all_stop_token_ids: BTreeSet::new(),

rust/src/llm/src/output.rs

+5 -3

@@ -69,7 +69,7 @@ pub enum FinishReason {

69 /// A retryable request-level internal error occurred.

70 Error,

71 /// A repetitive token pattern was detected.

72 - Repetition,

72 + Repetition(Option<StopReason>),

73 }

75 impl FinishReason {

@@ -87,7 +87,7 @@ impl FinishReason {

87 Self::Length => "length",

88 Self::Abort => "abort",

89 Self::Error => "error",

90 - Self::Repetition => "repetition",

90 + Self::Repetition(_) => "repetition",

91 }

92 }

@@ -96,6 +96,7 @@ impl FinishReason {

96 pub fn as_stop_reason(&self) -> Option<&StopReason> {

97 match self {

98 Self::Stop(stop_reason) => stop_reason.as_ref(),

99 + Self::Repetition(stop_reason) => stop_reason.as_ref(),

100 _ => None,

101 }

102 }

@@ -105,6 +106,7 @@ impl FinishReason {

106 pub fn into_stop_reason(self) -> Option<StopReason> {

107 match self {

108 Self::Stop(stop_reason) => stop_reason,

109 + Self::Repetition(stop_reason) => stop_reason,

110 _ => None,

111 }

112 }

@@ -119,7 +121,7 @@ fn finish_reason_from_engine(

121 EngineCoreFinishReason::Length => FinishReason::Length,

122 EngineCoreFinishReason::Abort => FinishReason::Abort,

123 EngineCoreFinishReason::Error => FinishReason::Error,

122 - EngineCoreFinishReason::Repetition => FinishReason::Repetition,

124 + EngineCoreFinishReason::Repetition => FinishReason::Repetition(stop_reason),

125 })

126 }

127

rust/src/server/src/grpc/convert.rs

+1 -1

@@ -332,7 +332,7 @@ fn to_finish_info(finished: &Finished, token_ids: &[u32]) -> pb::FinishInfo {

332 (PbFinishReason::Stop as i32, sr)

333 }

334 FinishReason::Length => (PbFinishReason::Length as i32, None),

335 - FinishReason::Abort | FinishReason::Error | FinishReason::Repetition => {

335 + FinishReason::Abort | FinishReason::Error | FinishReason::Repetition(_) => {

336 (PbFinishReason::Aborted as i32, None)

337 }

338 };

rust/src/server/src/routes/openai/chat_completions.rs

+1 -1

@@ -805,7 +805,7 @@ fn chat_finish_reason_to_openai(

805 FinishReason::Stop(_) => Ok("stop"),

806 FinishReason::Length => Ok("length"),

807 FinishReason::Abort => Ok("abort"),

808 - FinishReason::Repetition => Ok("stop"),

808 + FinishReason::Repetition(_) => Ok("repetition"),

809 FinishReason::Error => {

810 bail_server_error!("Internal server error");

811 }

rust/src/server/src/routes/openai/chat_completions/convert.rs

@@ -122,6 +122,7 @@ pub(super) fn prepare_chat_request(

122 frequency_penalty: request.frequency_penalty,

123 presence_penalty: request.presence_penalty,

124 repetition_penalty: request.repetition_penalty,

125 + repetition_detection: request.repetition_detection,

126 stop_token_ids: request.stop_token_ids,

127 ignore_eos: request.ignore_eos,

128 logit_bias: convert_logit_bias(request.logit_bias)?,

rust/src/server/src/routes/openai/chat_completions/types.rs

+2 -1

@@ -6,6 +6,7 @@ use serde_json::Value;

6 use serde_with::SerializeDisplay;

7 use validator::Validate;

8 use vllm_chat::ReasoningEffort;

9 +use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams;

11 use crate::routes::openai::utils::structured_outputs::ResponseFormat;

12 use crate::routes::openai::utils::types::{

@@ -236,7 +237,7 @@ pub struct ChatCompletionRequest {

237 pub vllm_xargs: Option<HashMap<String, Value>>,

238

239 /// Parameters for detecting repetitive N-gram patterns in output tokens

239 - pub repetition_detection: Option<Value>,

240 + pub repetition_detection: Option<RepetitionDetectionParams>,

241 }

242

243 impl Default for ChatCompletionRequest {

rust/src/server/src/routes/openai/chat_completions/validate.rs

-5

@@ -94,11 +94,6 @@ pub(super) fn validate_request_compat(

94 "mm_processor_kwargs",

95 "mm_processor_kwargs is not supported.",

96 )?;

97 - reject_non_default(

98 - request.repetition_detection.as_ref(),

99 - "repetition_detection",

100 - "repetition_detection is not supported.",

101 - )?;

98 Ok(())

99 }

rust/src/server/src/routes/openai/completions.rs

+32 -8

@@ -14,9 +14,11 @@ use axum::http::HeaderMap;

14 use axum::response::sse::{Event, Sse};

15 use axum::response::{IntoResponse, Response};

16 use futures::{Stream, StreamExt as _, pin_mut};

17 +use serde_json::Value;

18 use thiserror_ext::AsReport as _;

19 use tracing::{debug, error, info, trace};

20 use tracing_futures::Instrument as _;

21 +use vllm_engine_core_client::protocol::output::StopReason;

22 use vllm_text::{

23 DecodedPromptLogprobs, DecodedTextEvent, FinishReason, TextOutputStream,

24 TextOutputStreamExt as _,

@@ -142,9 +144,7 @@ async fn collect_completion(

144 .await

145 .map_err(|error| server_error!("completion stream failed: {}", error.to_report_string()))?;

146 let finish_reason = collected.finish_reason.clone();

145 - let stop_reason = finish_reason

146 - .as_stop_reason()

147 - .map(|sr| serde_json::to_value(sr).expect("StopReason must serialize to JSON"));

147 + let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);

148

149 let prompt_char_count = echo.as_ref().map(|prompt| text_len(prompt)).unwrap_or_default();

150 let logprobs = if requested_logprobs.is_some() && prompt_only {

@@ -181,7 +181,7 @@ async fn collect_completion(

181 Some(prompt) if prompt_only => prompt.clone(),

182 Some(prompt) => format!("{prompt}{}", collected.text),

183 };

184 - let finish_reason = completion_finish_reason_to_openai(finish_reason)?.to_string();

184 + let finish_reason = completion_finish_reason_to_openai(&finish_reason)?.to_string();

185 let usage = Usage::from_token_usage(collected.usage, enable_prompt_tokens_details);

186

187 if enable_log_requests {

@@ -441,21 +441,24 @@ fn final_chunk(

441 created: u64,

442 finish_reason: FinishReason,

443 ) -> Result<CompletionStreamResponse, ApiError> {

444 - let finish_reason = completion_finish_reason_to_openai(finish_reason)?;

444 + let stop_reason = finish_reason.as_stop_reason().map(stop_reason_to_json);

445 + let finish_reason = completion_finish_reason_to_openai(&finish_reason)?;

446

447 let mut chunk = CompletionStreamResponse::new(request_id, response_model, created);

448 chunk.choices.push(CompletionStreamChoice {

449 finish_reason: Some(finish_reason.to_string()),

450 + stop_reason,

451 ..Default::default()

452 });

453 Ok(chunk)

454 }

455

456 fn completion_finish_reason_to_openai(

455 - finish_reason: FinishReason,

457 + finish_reason: &FinishReason,

458 ) -> Result<&'static str, ApiError> {

459 match finish_reason {

458 - FinishReason::Stop(_) | FinishReason::Repetition => Ok("stop"),

460 + FinishReason::Stop(_) => Ok("stop"),

461 + FinishReason::Repetition(_) => Ok("repetition"),

462 FinishReason::Length => Ok("length"),

463 FinishReason::Abort => Ok("abort"),

464 FinishReason::Error => {

@@ -464,6 +467,10 @@ fn completion_finish_reason_to_openai(

467 }

468 }

469

470 +fn stop_reason_to_json(stop_reason: &StopReason) -> Value {

471 + serde_json::to_value(stop_reason).expect("StopReason must serialize to JSON")

472 +}

473 +

474 fn prompt_only_logprobs_to_openai(

475 prompt_logprobs: Option<&DecodedPromptLogprobs>,

476 prompt: &str,

@@ -577,6 +584,7 @@ fn done_sse_event() -> Event {

584 mod tests {

585 use futures::{StreamExt as _, stream};

586 use itertools::Itertools as _;

587 + use vllm_engine_core_client::protocol::output::StopReason;

588 use vllm_text::{

589 DecodedLogprobs, DecodedPositionLogprobs, DecodedPromptLogprobs, DecodedTextEvent,

590 DecodedTokenLogprob, FinishReason, Finished,

@@ -670,7 +678,9 @@ mod tests {

678 output_token_count: 2,

679 cached_token_count: 3,

680 },

673 - finish_reason: FinishReason::stop_eos(),

681 + finish_reason: FinishReason::Repetition(Some(StopReason::Text(

682 + "repetition_detected".to_string(),

683 + ))),

684 kv_transfer_params: None,

685 }),

686 }),

@@ -726,6 +736,20 @@ mod tests {

736 CompletionSseChunk::Usage(_) => panic!("expected regular chunk"),

737 }

738

739 + match &chunks[2] {

740 + CompletionSseChunk::Chunk(chunk) => {

741 + assert_eq!(

742 + chunk.choices[0].finish_reason.as_deref(),

743 + Some("repetition")

744 + );

745 + assert_eq!(

746 + chunk.choices[0].stop_reason,

747 + Some(serde_json::json!("repetition_detected"))

748 + );

749 + }

750 + CompletionSseChunk::Usage(_) => panic!("expected regular chunk"),

751 + }

752 +

753 match &chunks[3] {

754 CompletionSseChunk::Usage(chunk) => {

755 assert_eq!(

rust/src/server/src/routes/openai/completions/convert.rs

@@ -117,6 +117,7 @@ pub(super) fn prepare_completion_request(

117 frequency_penalty: request.frequency_penalty,

118 presence_penalty: request.presence_penalty,

119 repetition_penalty: request.repetition_penalty,

120 + repetition_detection: request.repetition_detection,

121 stop_token_ids: request.stop_token_ids,

122 ignore_eos: request.ignore_eos,

123 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;

3 use serde::{Deserialize, Serialize};

4 use serde_json::{Map, Value};

5 use validator::Validate;

6 +use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams;

7 use vllm_text::Prompt;

9 use crate::routes::openai::utils::types::{

@@ -100,6 +101,9 @@ pub struct CompletionRequest {

101 /// Repetition penalty for reducing repetitive text

102 pub repetition_penalty: Option<f32>,

103

104 + /// Parameters for detecting repetitive N-gram patterns in output tokens

105 + pub repetition_detection: Option<RepetitionDetectionParams>,

106 +

107 /// Length penalty for beam search

108 pub length_penalty: Option<f32>,

109

rust/src/text/src/error.rs

@@ -27,6 +27,8 @@ pub enum Error {

27 MinTokensExceedsMaxTokens { min_tokens: u32, max_tokens: u32 },

28 #[error("`thinking_token_budget` must be a non-negative integer or -1 for unlimited.")]

29 InvalidThinkingTokenBudget,

30 + #[error("invalid repetition detection params: {message}")]

31 + InvalidRepetitionDetection { message: String },

32 #[error("text request stream `{request_id}` closed before terminal output")]

33 StreamClosedBeforeTerminalOutput { request_id: String },

34 #[error(transparent)]

@@ -47,6 +49,7 @@ impl Error {

49 | Self::TokenIds(_)

50 | Self::MinTokensExceedsMaxTokens { .. }

51 | Self::InvalidThinkingTokenBudget

52 + | Self::InvalidRepetitionDetection { .. }

53 // An empty tokenized prompt detected later, at request prepare

54 // time, surfaces through the transparent Llm wrapper.

55 | Self::Llm(LlmError::EmptyPromptTokenIds { .. }) => true,

rust/src/text/src/lower.rs

+92 -1

@@ -5,7 +5,9 @@ pub(crate) mod token_ids;

6 use logprobs::validate_logprobs;

7 use token_ids::{validate_prompt_token_ids, validate_vocab_range};

8 -use vllm_engine_core_client::protocol::sampling::EngineCoreSamplingParams;

8 +use vllm_engine_core_client::protocol::sampling::{

9 + EngineCoreSamplingParams, RepetitionDetectionParams,

10 +};

11 use vllm_llm::GenerateRequest;

12 use vllm_tokenizer::Tokenizer;

@@ -93,6 +95,7 @@ pub fn lower_sampling_params(

95 frequency_penalty,

96 presence_penalty,

97 repetition_penalty,

98 + repetition_detection,

99 stop_token_ids,

100 ignore_eos,

101 logit_bias,

@@ -110,6 +113,7 @@ pub fn lower_sampling_params(

113 logprob_token_ids.as_deref(),

114 sampling_limits,

115 )?;

116 + validate_repetition_detection(repetition_detection.as_ref())?;

117

118 // Mirrors the model-generation-config inheritance used by vLLM's OpenAI chat

119 // 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(

167 frequency_penalty,

168 presence_penalty,

169 repetition_penalty,

170 + repetition_detection: repetition_detection.filter(|p| !p.is_disabled()),

171 stop_token_ids,

172 eos_token_id: (!ignore_eos).then_some(primary_eos_token_id).flatten(),

173 all_stop_token_ids,

@@ -194,6 +199,33 @@ fn normalize_thinking_token_budget(value: Option<i64>) -> Result<Option<u64>> {

199 }

200 }

201

202 +fn validate_repetition_detection(params: Option<&RepetitionDetectionParams>) -> Result<()> {

203 + let Some(params) = params else {

204 + return Ok(());

205 + };

206 +

207 + if params.min_pattern_size > params.max_pattern_size {

208 + return Err(Error::InvalidRepetitionDetection {

209 + message: format!(

210 + "`min_pattern_size` must be less than or equal to \

211 + `max_pattern_size`, got min_pattern_size={}, \

212 + max_pattern_size={}",

213 + params.min_pattern_size, params.max_pattern_size

214 + ),

215 + });

216 + }

217 + if params.max_pattern_size > 0 && params.min_count < 2 {

218 + return Err(Error::InvalidRepetitionDetection {

219 + message: format!(

220 + "`min_count` must be at least 2, got min_count={}",

221 + params.min_count

222 + ),

223 + });

224 + }

225 +

226 + Ok(())

227 +}

228 +

229 /// Convert bad-word strings into token-ID sequences, following the Python vLLM

230 /// logic in `SamplingParams.update_from_tokenizer()`.

231 ///

@@ -390,6 +422,59 @@ mod tests {

422 ));

423 }

424

425 + #[test]

426 + fn lower_sampling_params_validates_repetition_detection() {

427 + let lower = |repetition_detection| {

428 + lower_sampling_params_with_limits(

429 + SamplingParams {

430 + repetition_detection,

431 + ..SamplingParams::default()

432 + },

433 + sample_sampling_limits(),

434 + )

435 + };

436 +

437 + let enabled = RepetitionDetectionParams {

438 + max_pattern_size: 4,

439 + min_pattern_size: 2,

440 + min_count: 2,

441 + };

442 + assert_eq!(

443 + lower(Some(enabled.clone())).unwrap().repetition_detection,

444 + Some(enabled)

445 + );

446 +

447 + let disabled = RepetitionDetectionParams {

448 + max_pattern_size: 0,

449 + min_pattern_size: 0,

450 + min_count: 0,

451 + };

452 + assert_eq!(lower(Some(disabled)).unwrap().repetition_detection, None);

453 +

454 + let error = lower(Some(RepetitionDetectionParams {

455 + max_pattern_size: 1,

456 + min_pattern_size: 2,

457 + min_count: 2,

458 + }))

459 + .unwrap_err();

460 + let Error::InvalidRepetitionDetection { message } = error else {

461 + panic!("expected repetition_detection validation error");

462 + };

463 + assert!(message.contains("min_pattern_size=2"));

464 + assert!(message.contains("max_pattern_size=1"));

465 +

466 + let error = lower(Some(RepetitionDetectionParams {

467 + max_pattern_size: 1,

468 + min_pattern_size: 1,

469 + min_count: 1,

470 + }))

471 + .unwrap_err();

472 + let Error::InvalidRepetitionDetection { message } = error else {

473 + panic!("expected repetition_detection validation error");

474 + };

475 + assert!(message.contains("min_count=1"));

476 + }

477 +

478 #[test]

479 fn lower_text_request_applies_python_style_eos_hints() {

480 let prepared = lower_text_request(

@@ -417,6 +502,7 @@ mod tests {

502 frequency_penalty: 0.0,

503 presence_penalty: 0.0,

504 repetition_penalty: 1.0,

505 + repetition_detection: None,

506 stop_token_ids: [

507 77,

508 ],

@@ -469,6 +555,7 @@ mod tests {

555 frequency_penalty: 0.0,

556 presence_penalty: 0.0,

557 repetition_penalty: 1.0,

558 + repetition_detection: None,

559 stop_token_ids: [],

560 eos_token_id: None,

561 all_stop_token_ids: {

@@ -600,6 +687,7 @@ mod tests {

687 frequency_penalty: 0.0,

688 presence_penalty: 0.0,

689 repetition_penalty: 1.0,

690 + repetition_detection: None,

691 stop_token_ids: [

692 151643,

693 ],

@@ -662,6 +750,7 @@ mod tests {

750 frequency_penalty: 0.0,

751 presence_penalty: 0.0,

752 repetition_penalty: 1.0,

753 + repetition_detection: None,

754 stop_token_ids: [

755 11,

756 77,

@@ -732,6 +821,7 @@ mod tests {

821 frequency_penalty: 0.0,

822 presence_penalty: 0.0,

823 repetition_penalty: 1.2,

824 + repetition_detection: None,

825 stop_token_ids: [],

826 eos_token_id: None,

827 all_stop_token_ids: {},

@@ -980,6 +1070,7 @@ mod tests {

1070 frequency_penalty: 0.0,

1071 presence_penalty: 0.0,

1072 repetition_penalty: 1.2,

1073 + repetition_detection: None,

1074 stop_token_ids: [],

1075 eos_token_id: None,

1076 all_stop_token_ids: {},

rust/src/text/src/request.rs

@@ -6,6 +6,7 @@ use serde_json::Value;

6 use vllm_engine_core_client::protocol::lora::LoraRequest;

7 use vllm_engine_core_client::protocol::multimodal::MmFeatures;

8 use vllm_engine_core_client::protocol::request::ReasoningParserKwargs;

9 +use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams;

10 use vllm_engine_core_client::protocol::structured_outputs::StructuredOutputsParams;

12 use crate::error::{Error, Result};

@@ -83,6 +84,9 @@ pub struct SamplingParams {

84 /// Repetition penalty applied by the sampler. `None` means no explicit user

85 /// override.

86 pub repetition_penalty: Option<f32>,

87 + /// Parameters for detecting repetitive N-gram patterns. `None` means no

88 + /// explicit user override.

89 + pub repetition_detection: Option<RepetitionDetectionParams>,

90 /// Explicit stop token IDs provided by the caller. `None` means no explicit

91 /// user override.

92 pub stop_token_ids: Option<Vec<u32>>,

@@ -130,6 +134,7 @@ impl Default for SamplingParams {

134 frequency_penalty: None,

135 presence_penalty: None,

136 repetition_penalty: None,

137 + repetition_detection: None,

138 stop_token_ids: None,

139 ignore_eos: false,

140 logit_bias: None,