vllm_cpp/
params.rs

1use std::collections::HashMap;
2use std::ffi::CString;
3use std::fmt;
4use std::mem::{align_of, size_of};
5use std::os::raw::{c_char, c_void};
6use std::panic::{catch_unwind, AssertUnwindSafe};
7use std::ptr;
8use std::slice;
9use std::sync::atomic::{AtomicU8, AtomicUsize, Ordering};
10use std::sync::{Arc, Mutex, OnceLock, Weak};
11use std::thread::{self, ThreadId};
12
13use vllm_cpp_sys as ffi;
14
15use crate::error::{invalid_configuration, Error};
16
17const NATIVE_DEFAULT_MAX_TOKENS: u32 = 16;
18
19/// Native scheduler admission order.
20///
21/// Raw and serde chat request JSON can carry a `priority` field that the native
22/// OpenAI-compatible path parses and submits. Direct completion, completion
23/// streaming, and [`crate::Request`] submissions currently default to priority zero
24/// and tie by arrival; caller-selected priorities for those direct APIs require a
25/// future C ABI/API change.
26#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
27pub enum SchedulerPolicy {
28    /// Process requests in arrival order.
29    #[default]
30    Fcfs,
31    /// Order requests by priority and then arrival time.
32    ///
33    /// This variant selects the native priority queue; it does not itself assign a
34    /// priority to a request.
35    Priority,
36
37    /// Prefer requests sharing the longest cached prefix.
38    LongestPrefixMatch,
39}
40
41impl SchedulerPolicy {
42    pub(crate) const fn as_str(self) -> &'static str {
43        match self {
44            Self::Fcfs => "fcfs",
45            Self::Priority => "priority",
46            Self::LongestPrefixMatch => "lpm",
47        }
48    }
49}
50
51/// A native tri-state setting whose default is resolved by vllm.cpp.
52#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)]
53pub enum Toggle {
54    /// Let vllm.cpp resolve the model or environment default.
55    #[default]
56    Default,
57    /// Force the feature on.
58    On,
59    /// Force the feature off.
60    Off,
61}
62
63impl Toggle {
64    pub(crate) const fn as_native(self) -> i32 {
65        match self {
66            Self::Default => 0,
67            Self::On => 1,
68            Self::Off => 2,
69        }
70    }
71}
72
73type DynLogitsProcessor = dyn Fn(&[i32], &mut [f32]) + Send + Sync + 'static;
74
75#[derive(Clone)]
76struct LogitsProcessor {
77    callback: Arc<DynLogitsProcessor>,
78}
79
80impl fmt::Debug for LogitsProcessor {
81    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
82        formatter.write_str("LogitsProcessor { .. }")
83    }
84}
85
86impl PartialEq for LogitsProcessor {
87    fn eq(&self, other: &Self) -> bool {
88        Arc::ptr_eq(&self.callback, &other.callback)
89    }
90}
91
92/// One engine-side structured decoding constraint.
93#[derive(Clone, Debug, Eq, PartialEq)]
94pub enum StructuredOutput {
95    JsonSchema(String),
96    Regex(String),
97    Choice(Vec<String>),
98    Grammar(String),
99    JsonObject,
100}
101
102/// Owned sampling configuration for one generation request.
103#[derive(Clone, Debug, PartialEq)]
104pub struct SamplingParams {
105    temperature: f32,
106    top_p: f32,
107    top_k: i32,
108    min_p: f32,
109    max_tokens: Option<u32>,
110    seed: Option<u64>,
111    presence_penalty: f32,
112    frequency_penalty: f32,
113    repetition_penalty: f32,
114    min_tokens: u32,
115    ignore_eos: bool,
116    stop: Vec<String>,
117    structured_output: Option<StructuredOutput>,
118    logits_processor: Option<LogitsProcessor>,
119}
120
121impl Default for SamplingParams {
122    fn default() -> Self {
123        Self {
124            temperature: 1.0,
125            top_p: 1.0,
126            top_k: 0,
127            min_p: 0.0,
128            max_tokens: Some(NATIVE_DEFAULT_MAX_TOKENS),
129            seed: None,
130            presence_penalty: 0.0,
131            frequency_penalty: 0.0,
132            repetition_penalty: 1.0,
133            min_tokens: 0,
134            ignore_eos: false,
135            stop: Vec::new(),
136            structured_output: None,
137            logits_processor: None,
138        }
139    }
140}
141
142impl SamplingParams {
143    /// Returns deterministic argmax sampling with native defaults otherwise.
144    #[must_use]
145    pub fn greedy() -> Self {
146        Self::default().temperature(0.0)
147    }
148
149    #[must_use]
150    pub fn temperature(mut self, value: f32) -> Self {
151        self.temperature = value;
152        self
153    }
154
155    #[must_use]
156    pub fn top_p(mut self, value: f32) -> Self {
157        self.top_p = value;
158        self
159    }
160
161    #[must_use]
162    pub fn top_k(mut self, value: i32) -> Self {
163        self.top_k = value;
164        self
165    }
166
167    #[must_use]
168    pub fn min_p(mut self, value: f32) -> Self {
169        self.min_p = value;
170        self
171    }
172
173    /// Sets a finite generation limit.
174    ///
175    /// Zero is invalid; use [`unbounded`](Self::unbounded) to request native
176    /// unbounded generation explicitly.
177    #[must_use]
178    pub fn max_tokens(mut self, value: u32) -> Self {
179        self.max_tokens = Some(value);
180        self
181    }
182
183    /// Removes the generation limit.
184    #[must_use]
185    pub fn unbounded(mut self) -> Self {
186        self.max_tokens = None;
187        self
188    }
189
190    #[must_use]
191    pub fn seed(mut self, value: u64) -> Self {
192        self.seed = Some(value);
193        self
194    }
195
196    #[must_use]
197    pub fn clear_seed(mut self) -> Self {
198        self.seed = None;
199        self
200    }
201
202    #[must_use]
203    pub fn presence_penalty(mut self, value: f32) -> Self {
204        self.presence_penalty = value;
205        self
206    }
207
208    #[must_use]
209    pub fn frequency_penalty(mut self, value: f32) -> Self {
210        self.frequency_penalty = value;
211        self
212    }
213
214    #[must_use]
215    pub fn repetition_penalty(mut self, value: f32) -> Self {
216        self.repetition_penalty = value;
217        self
218    }
219
220    #[must_use]
221    pub fn min_tokens(mut self, value: u32) -> Self {
222        self.min_tokens = value;
223        self
224    }
225
226    #[must_use]
227    pub fn ignore_eos(mut self, value: bool) -> Self {
228        self.ignore_eos = value;
229        self
230    }
231
232    #[must_use]
233    pub fn stop(mut self, value: impl Into<String>) -> Self {
234        self.stop.push(value.into());
235        self
236    }
237
238    #[must_use]
239    pub fn stop_all<I, S>(mut self, values: I) -> Self
240    where
241        I: IntoIterator<Item = S>,
242        S: Into<String>,
243    {
244        self.stop.extend(values.into_iter().map(Into::into));
245        self
246    }
247
248    #[must_use]
249    pub fn structured_output(mut self, value: StructuredOutput) -> Self {
250        self.structured_output = Some(value);
251        self
252    }
253
254    /// Installs a host-side processor that can inspect generated token IDs and
255    /// mutate one request's logits before each sampling step.
256    ///
257    /// The processor may run concurrently for different requests, so it must be
258    /// `Send + Sync`. Cloned parameters share the processor. A panic is contained
259    /// before the C boundary and reported as [`Error::LogitsProcessorPanicked`]
260    /// after the bounded generation call or from [`crate::Request::wait`]. The
261    /// callback state remains registered through the call or request lifetime;
262    /// stale native invocations after cleanup become no-ops.
263    #[must_use]
264    pub fn logits_processor<F>(mut self, processor: F) -> Self
265    where
266        F: Fn(&[i32], &mut [f32]) + Send + Sync + 'static,
267    {
268        self.logits_processor = Some(LogitsProcessor {
269            callback: Arc::new(processor),
270        });
271        self
272    }
273
274    /// Removes a previously configured custom logits processor.
275    #[must_use]
276    pub fn clear_logits_processor(mut self) -> Self {
277        self.logits_processor = None;
278        self
279    }
280
281    pub(crate) fn marshal(&self) -> Result<MarshaledSamplingParams, Error> {
282        MarshaledSamplingParams::new(self)
283    }
284}
285
286pub(crate) struct MarshaledSamplingParams {
287    raw: ffi::vllm_sampling_params,
288    _stop: Vec<CString>,
289    _stop_pointers: Vec<*const c_char>,
290    _structured_string: Option<CString>,
291    _choices: Vec<CString>,
292    _choice_pointers: Vec<*const c_char>,
293    logits_processor: Option<LogitsProcessorRegistration>,
294}
295
296impl MarshaledSamplingParams {
297    fn new(params: &SamplingParams) -> Result<Self, Error> {
298        // ABI equality is checked before this struct-returning call.
299        let mut raw = unsafe { ffi::vllm_sampling_params_default() };
300        raw.temperature = params.temperature;
301        raw.top_p = params.top_p;
302        raw.top_k = params.top_k;
303        raw.min_p = params.min_p;
304        raw.max_tokens = optional_u32_to_i32(params.max_tokens, "max_tokens")?;
305        raw.seed = params.seed.unwrap_or(0);
306        raw.has_seed = i32::from(params.seed.is_some());
307        raw.presence_penalty = params.presence_penalty;
308        raw.frequency_penalty = params.frequency_penalty;
309        raw.repetition_penalty = params.repetition_penalty;
310        raw.min_tokens = u32_to_i32(params.min_tokens, "min_tokens")?;
311        raw.ignore_eos = i32::from(params.ignore_eos);
312
313        let stop = strings_to_cstrings(&params.stop, "stop string")?;
314        let stop_pointers = stop.iter().map(|value| value.as_ptr()).collect::<Vec<_>>();
315        raw.stop = pointer_or_null(&stop_pointers);
316        raw.n_stop = length_to_i32(stop_pointers.len(), "stop strings")?;
317
318        let mut structured_string = None;
319        let mut choices = Vec::new();
320        let mut choice_pointers = Vec::new();
321        if let Some(structured) = &params.structured_output {
322            match structured {
323                StructuredOutput::JsonSchema(value) => {
324                    structured_string = Some(to_cstring(value, "JSON schema")?);
325                    raw.structured_json = structured_string.as_ref().unwrap().as_ptr();
326                }
327                StructuredOutput::Regex(value) => {
328                    structured_string = Some(to_cstring(value, "structured regex")?);
329                    raw.structured_regex = structured_string.as_ref().unwrap().as_ptr();
330                }
331                StructuredOutput::Choice(values) => {
332                    if values.is_empty() {
333                        return Err(invalid_configuration(
334                            "structured choices must contain at least one value",
335                        ));
336                    }
337                    choices = strings_to_cstrings(values, "structured choice")?;
338                    choice_pointers = choices
339                        .iter()
340                        .map(|value| value.as_ptr())
341                        .collect::<Vec<_>>();
342                    raw.structured_choice = pointer_or_null(&choice_pointers);
343                    raw.n_structured_choice =
344                        length_to_i32(choice_pointers.len(), "structured choices")?;
345                }
346                StructuredOutput::Grammar(value) => {
347                    structured_string = Some(to_cstring(value, "structured grammar")?);
348                    raw.structured_grammar = structured_string.as_ref().unwrap().as_ptr();
349                }
350                StructuredOutput::JsonObject => raw.structured_json_object = 1,
351            }
352        }
353
354        let mut logits_processor = None;
355        if let Some(processor) = &params.logits_processor {
356            if params.max_tokens.is_none() || params.max_tokens == Some(0) {
357                return Err(invalid_configuration(
358                    "custom logits processors require bounded max_tokens because the native callback cannot abort generation",
359                ));
360            }
361            let registration = LogitsProcessorRegistration::new(Arc::clone(&processor.callback));
362            raw.logits_processor = Some(logits_processor_trampoline);
363            raw.logits_processor_user_data = registration.user_data();
364            logits_processor = Some(registration);
365        }
366
367        Ok(Self {
368            raw,
369            _stop: stop,
370            _stop_pointers: stop_pointers,
371            _structured_string: structured_string,
372            _choices: choices,
373            _choice_pointers: choice_pointers,
374            logits_processor,
375        })
376    }
377
378    pub(crate) const fn raw(&self) -> &ffi::vllm_sampling_params {
379        &self.raw
380    }
381
382    pub(crate) fn logits_processor_error(&self) -> Option<Error> {
383        self.logits_processor
384            .as_ref()
385            .and_then(LogitsProcessorRegistration::error)
386    }
387
388    pub(crate) fn take_logits_processor(&mut self) -> Option<LogitsProcessorRegistration> {
389        self.logits_processor.take()
390    }
391}
392
393const PROCESSOR_OK: u8 = 0;
394const PROCESSOR_PANICKED: u8 = 1;
395const PROCESSOR_INVALID_INPUT: u8 = 2;
396
397static NEXT_PROCESSOR_ID: AtomicUsize = AtomicUsize::new(1);
398static LOGITS_PROCESSORS: OnceLock<Mutex<HashMap<usize, Weak<LogitsProcessorState>>>> =
399    OnceLock::new();
400
401pub(crate) struct LogitsProcessorRegistration {
402    id: usize,
403    state: Arc<LogitsProcessorState>,
404}
405
406impl LogitsProcessorRegistration {
407    fn new(callback: Arc<DynLogitsProcessor>) -> Self {
408        let id = NEXT_PROCESSOR_ID
409            .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |id| id.checked_add(1))
410            .unwrap_or_else(|_| std::process::abort());
411        let state = Arc::new(LogitsProcessorState::new(callback));
412        lock_unpoisoned(logits_processor_registry()).insert(id, Arc::downgrade(&state));
413        Self { id, state }
414    }
415
416    fn user_data(&self) -> *mut c_void {
417        ptr::without_provenance_mut(self.id)
418    }
419
420    pub(crate) fn error(&self) -> Option<Error> {
421        self.state.error()
422    }
423
424    pub(crate) fn is_active_on_current_thread(&self) -> bool {
425        self.state.is_active_on_current_thread()
426    }
427}
428
429impl Drop for LogitsProcessorRegistration {
430    fn drop(&mut self) {
431        lock_unpoisoned(logits_processor_registry()).remove(&self.id);
432    }
433}
434
435fn logits_processor_registry() -> &'static Mutex<HashMap<usize, Weak<LogitsProcessorState>>> {
436    LOGITS_PROCESSORS.get_or_init(|| Mutex::new(HashMap::new()))
437}
438
439fn registered_logits_processor(id: usize) -> Option<Arc<LogitsProcessorState>> {
440    lock_unpoisoned(logits_processor_registry())
441        .get(&id)
442        .and_then(Weak::upgrade)
443}
444
445struct LogitsProcessorState {
446    callback: Arc<DynLogitsProcessor>,
447    failure: AtomicU8,
448    active_threads: Mutex<Vec<ThreadId>>,
449}
450
451impl LogitsProcessorState {
452    fn new(callback: Arc<DynLogitsProcessor>) -> Self {
453        Self {
454            callback,
455            failure: AtomicU8::new(PROCESSOR_OK),
456            active_threads: Mutex::new(Vec::new()),
457        }
458    }
459
460    pub(crate) fn error(&self) -> Option<Error> {
461        match self.failure.load(Ordering::Acquire) {
462            PROCESSOR_OK => None,
463            PROCESSOR_PANICKED => Some(Error::LogitsProcessorPanicked),
464            PROCESSOR_INVALID_INPUT => Some(Error::Runtime {
465                message: "native logits processor callback received invalid pointers or lengths"
466                    .to_owned(),
467            }),
468            _ => Some(Error::Runtime {
469                message: "native logits processor callback entered an unknown failure state"
470                    .to_owned(),
471            }),
472        }
473    }
474
475    pub(crate) fn is_active_on_current_thread(&self) -> bool {
476        let current = thread::current().id();
477        lock_unpoisoned(&self.active_threads).contains(&current)
478    }
479
480    fn record_failure(&self, failure: u8) {
481        let _ = self.failure.compare_exchange(
482            PROCESSOR_OK,
483            failure,
484            Ordering::AcqRel,
485            Ordering::Acquire,
486        );
487    }
488}
489
490struct ActiveProcessorGuard<'state> {
491    state: &'state LogitsProcessorState,
492    thread_id: ThreadId,
493}
494
495impl<'state> ActiveProcessorGuard<'state> {
496    fn enter(state: &'state LogitsProcessorState) -> Self {
497        let thread_id = thread::current().id();
498        lock_unpoisoned(&state.active_threads).push(thread_id);
499        Self { state, thread_id }
500    }
501}
502
503impl Drop for ActiveProcessorGuard<'_> {
504    fn drop(&mut self) {
505        let mut active = lock_unpoisoned(&self.state.active_threads);
506        if let Some(index) = active.iter().rposition(|id| *id == self.thread_id) {
507            active.swap_remove(index);
508        }
509    }
510}
511
512unsafe extern "C" fn logits_processor_trampoline(
513    token_ids: *const i32,
514    n_token_ids: i32,
515    logits: *mut f32,
516    vocab_size: i32,
517    user_data: *mut c_void,
518) {
519    if user_data.is_null() {
520        return;
521    }
522    let Some(state) = registered_logits_processor(user_data.addr()) else {
523        return;
524    };
525    if state.failure.load(Ordering::Acquire) != PROCESSOR_OK {
526        return;
527    }
528    if n_token_ids < 0
529        || vocab_size <= 0
530        || (n_token_ids > 0 && token_ids.is_null())
531        || logits.is_null()
532        || (n_token_ids > 0 && !valid_slice_layout(token_ids, n_token_ids as usize))
533        || !valid_slice_layout(logits, vocab_size as usize)
534    {
535        state.record_failure(PROCESSOR_INVALID_INPUT);
536        return;
537    }
538
539    let _active = ActiveProcessorGuard::enter(&state);
540    let tokens = if n_token_ids == 0 {
541        &[]
542    } else {
543        // SAFETY: the native callback contract lends this aligned token slice for
544        // this invocation, and the validated length fits Rust slice bounds.
545        unsafe { slice::from_raw_parts(token_ids, n_token_ids as usize) }
546    };
547    // SAFETY: the native callback contract lends this aligned, uniquely mutable
548    // logits row for this invocation, and the validated length fits slice bounds.
549    let logits = unsafe { slice::from_raw_parts_mut(logits, vocab_size as usize) };
550    if let Err(payload) = catch_unwind(AssertUnwindSafe(|| (state.callback)(tokens, logits))) {
551        state.record_failure(PROCESSOR_PANICKED);
552        discard_panic_payload(payload);
553    }
554}
555
556fn valid_slice_layout<T>(pointer: *const T, length: usize) -> bool {
557    !pointer.is_null()
558        && (pointer as usize) % align_of::<T>() == 0
559        && length <= (isize::MAX as usize) / size_of::<T>()
560}
561
562fn discard_panic_payload(payload: Box<dyn std::any::Any + Send>) {
563    if let Err(payload) = catch_unwind(AssertUnwindSafe(|| drop(payload))) {
564        std::mem::forget(payload);
565    }
566}
567
568fn lock_unpoisoned<T>(mutex: &Mutex<T>) -> std::sync::MutexGuard<'_, T> {
569    mutex
570        .lock()
571        .unwrap_or_else(std::sync::PoisonError::into_inner)
572}
573
574pub(crate) fn to_cstring(value: &str, field: &'static str) -> Result<CString, Error> {
575    CString::new(value).map_err(|_| Error::InteriorNul { field })
576}
577
578fn strings_to_cstrings(values: &[String], field: &'static str) -> Result<Vec<CString>, Error> {
579    values
580        .iter()
581        .map(|value| to_cstring(value, field))
582        .collect()
583}
584
585fn pointer_or_null(values: &[*const c_char]) -> *const *const c_char {
586    if values.is_empty() {
587        ptr::null()
588    } else {
589        values.as_ptr()
590    }
591}
592
593fn optional_u32_to_i32(value: Option<u32>, field: &'static str) -> Result<i32, Error> {
594    match value {
595        Some(0) => Err(invalid_configuration(format!(
596            "{field} must be greater than zero; use unbounded() for no limit"
597        ))),
598        Some(value) => u32_to_i32(value, field),
599        None => Ok(0),
600    }
601}
602
603fn u32_to_i32(value: u32, field: &'static str) -> Result<i32, Error> {
604    i32::try_from(value)
605        .map_err(|_| invalid_configuration(format!("{field} exceeds native i32 range")))
606}
607
608fn length_to_i32(value: usize, field: &'static str) -> Result<i32, Error> {
609    i32::try_from(value).map_err(|_| invalid_configuration(format!("too many {field}")))
610}
611
612#[cfg(test)]
613mod tests {
614    use super::{logits_processor_trampoline, SamplingParams};
615    use crate::Error;
616    use std::sync::atomic::{AtomicUsize, Ordering};
617    use std::sync::Arc;
618
619    #[test]
620    fn marshals_and_invokes_custom_logits_processor() {
621        let params = SamplingParams::default()
622            .max_tokens(2)
623            .logits_processor(|tokens, logits| {
624                assert_eq!(tokens, &[3, 5]);
625                logits[1] = 9.0;
626            });
627        let marshaled = params.marshal().expect("marshal processor");
628        let callback = marshaled
629            .raw()
630            .logits_processor
631            .expect("processor callback");
632        let mut logits = [1.0, 2.0, 3.0];
633        let tokens = [3, 5];
634        unsafe {
635            callback(
636                tokens.as_ptr(),
637                tokens.len() as i32,
638                logits.as_mut_ptr(),
639                logits.len() as i32,
640                marshaled.raw().logits_processor_user_data,
641            );
642        }
643        assert_eq!(logits, [1.0, 9.0, 3.0]);
644        assert_eq!(marshaled.logits_processor_error(), None);
645    }
646
647    #[test]
648    fn stale_processor_user_data_is_a_noop() {
649        let calls = Arc::new(AtomicUsize::new(0));
650        let user_data = {
651            let calls = Arc::clone(&calls);
652            let params = SamplingParams::default()
653                .max_tokens(1)
654                .logits_processor(move |_, _| {
655                    calls.fetch_add(1, Ordering::Relaxed);
656                });
657            let marshaled = params.marshal().expect("marshal processor");
658            marshaled.raw().logits_processor_user_data
659        };
660        let mut logits = [1.0];
661        unsafe {
662            logits_processor_trampoline(
663                std::ptr::null(),
664                0,
665                logits.as_mut_ptr(),
666                logits.len() as i32,
667                user_data,
668            );
669        }
670        assert_eq!(calls.load(Ordering::Relaxed), 0);
671    }
672
673    #[test]
674    fn contains_processor_panic_and_skips_later_calls() {
675        let params = SamplingParams::default()
676            .max_tokens(2)
677            .logits_processor(|_, _| panic!("processor panic"));
678        let marshaled = params.marshal().expect("marshal processor");
679        let mut logits = [1.0, 2.0];
680        unsafe {
681            logits_processor_trampoline(
682                std::ptr::null(),
683                0,
684                logits.as_mut_ptr(),
685                logits.len() as i32,
686                marshaled.raw().logits_processor_user_data,
687            );
688        }
689        assert_eq!(
690            marshaled.logits_processor_error(),
691            Some(Error::LogitsProcessorPanicked)
692        );
693        unsafe {
694            logits_processor_trampoline(
695                std::ptr::null(),
696                0,
697                logits.as_mut_ptr(),
698                logits.len() as i32,
699                marshaled.raw().logits_processor_user_data,
700            );
701        }
702    }
703
704    #[test]
705    fn rejects_zero_or_unbounded_processor_and_invalid_native_shape() {
706        for params in [
707            SamplingParams::default().max_tokens(0),
708            SamplingParams::default()
709                .unbounded()
710                .logits_processor(|_, _| {}),
711        ] {
712            let error = params.marshal().err().expect("invalid bounds rejection");
713            assert!(matches!(error, Error::InvalidConfiguration { .. }));
714        }
715
716        let params = SamplingParams::default().logits_processor(|_, _| {});
717        let marshaled = params.marshal().expect("marshal processor");
718        unsafe {
719            logits_processor_trampoline(
720                std::ptr::null(),
721                -1,
722                std::ptr::null_mut(),
723                0,
724                marshaled.raw().logits_processor_user_data,
725            );
726        }
727        assert!(matches!(
728            marshaled.logits_processor_error(),
729            Some(Error::Runtime { .. })
730        ));
731    }
732}
vllm_cpp/params.rs

vllm_cpp/
params.rs
