diff --git a/vllm-cpp/src/abi.rs b/vllm-cpp/src/abi.rs new file mode 100644 index 0000000..d8cf2ec --- /dev/null +++ b/vllm-cpp/src/abi.rs @@ -0,0 +1,109 @@ +use vllm_cpp_sys as ffi; + +use crate::Error; + +/// Proof that the linked library exactly matches the generated ABI. +pub(crate) struct Compatibility { + _private: (), +} + +impl Compatibility { + pub(crate) fn check() -> Result { + // SAFETY: this base ABI function takes no pointers or versioned structs. + let actual = unsafe { ffi::vllm_abi_version() }; + Self::from_actual(actual) + } + + pub(crate) fn model_params_default(&self) -> ffi::vllm_model_params { + // SAFETY: possession of this token proves exact ABI equality for this + // engine construction before returning the versioned struct by value. + unsafe { ffi::vllm_model_params_default() } + } + + pub(crate) fn sampling_params_default(&self) -> ffi::vllm_sampling_params { + // SAFETY: the engine retained this token after exact ABI equality was + // established, so this versioned struct may be returned by value. + unsafe { ffi::vllm_sampling_params_default() } + } + + fn from_actual(actual: i32) -> Result { + let expected = ffi::VLLM_ABI_VERSION as i32; + if actual != expected { + return Err(Error::AbiMismatch { expected, actual }); + } + Ok(Self { _private: () }) + } + + #[cfg(test)] + pub(crate) fn check_with(abi_version: impl FnOnce() -> i32) -> Result { + Self::from_actual(abi_version()) + } + + #[cfg(test)] + fn model_params_default_with( + &self, + default: impl FnOnce() -> ffi::vllm_model_params, + ) -> ffi::vllm_model_params { + default() + } + + #[cfg(test)] + fn sampling_params_default_with( + &self, + default: impl FnOnce() -> ffi::vllm_sampling_params, + ) -> ffi::vllm_sampling_params { + default() + } +} + +#[cfg(test)] +mod tests { + use std::cell::RefCell; + + use super::Compatibility; + use crate::Error; + + #[test] + fn mismatch_produces_no_token_or_default_access() { + let calls = RefCell::new(Vec::new()); + let result = Compatibility::check_with(|| { + calls.borrow_mut().push("abi"); + 10 + }); + + assert!(matches!( + result, + Err(Error::AbiMismatch { + expected: 17, + actual: 10 + }) + )); + assert_eq!(*calls.borrow(), ["abi"]); + } + + #[test] + fn compatibility_precedes_both_by_value_defaults() { + let calls = RefCell::new(Vec::new()); + let compatibility = Compatibility::check_with(|| { + calls.borrow_mut().push("abi"); + 17 + }) + .expect("matching compatibility token"); + + compatibility.model_params_default_with(|| { + calls.borrow_mut().push("model_default"); + // SAFETY: every field in this generated C struct permits zero. + unsafe { std::mem::zeroed() } + }); + compatibility.sampling_params_default_with(|| { + calls.borrow_mut().push("sampling_default"); + // SAFETY: every field in this generated C struct permits zero. + unsafe { std::mem::zeroed() } + }); + + assert_eq!( + *calls.borrow(), + ["abi", "model_default", "sampling_default"] + ); + } +} diff --git a/vllm-cpp/src/engine.rs b/vllm-cpp/src/engine.rs index 4ab3d90..cafde32 100644 --- a/vllm-cpp/src/engine.rs +++ b/vllm-cpp/src/engine.rs @@ -1,26 +1,64 @@ use std::ffi::{CStr, CString}; +use std::marker::PhantomData; use std::mem::MaybeUninit; use std::os::raw::c_char; use std::path::{Path, PathBuf}; use std::ptr::{self, NonNull}; +use std::rc::Rc; use std::sync::Arc; use vllm_cpp_sys as ffi; +use crate::abi::Compatibility; use crate::callback::{ callback_trampoline, CallbackState, StreamControl, StreamEvent, StreamOutcome, }; use crate::error::{invalid_configuration, status_result, Error}; -use crate::params::{SamplingParams, SchedulerPolicy, Toggle}; +use crate::params::{Device, SamplingParams, SchedulerPolicy, Toggle}; -/// A cloneable vllm.cpp serving engine. +/// A cloneable, `Send + Sync` vllm.cpp text-generation engine. #[derive(Clone)] pub struct Engine { pub(crate) inner: Arc, } -pub(crate) struct EngineInner { +pub(crate) struct TextTask; +struct TranscriptionTask; +struct EmbeddingTask; + +pub(crate) struct OwnedEngine { pub(crate) raw: NonNull, + pub(crate) compatibility: Compatibility, + _task: PhantomData, + _not_send_sync: PhantomData>, +} + +struct LoadedEngine { + raw: NonNull, + compatibility: Compatibility, + _task: PhantomData, +} + +pub(crate) type EngineInner = OwnedEngine; + +/// A thread-local RAII owner for a native transcription-task engine. +/// +/// ABI 17 has no task query, so loading cannot prove that a checkpoint supports +/// transcription. Native task selection and future wrong-task diagnostics remain +/// authoritative. This owner intentionally exposes no operation yet and is +/// neither `Send` nor `Sync`. +pub struct TranscriptionEngine { + _inner: OwnedEngine, +} + +/// A thread-local RAII owner for a native embedding-task engine. +/// +/// ABI 17 has no task query, so loading cannot prove that a checkpoint supports +/// embeddings. Native task selection and future wrong-task diagnostics remain +/// authoritative. This owner intentionally exposes no operation yet and is +/// neither `Send` nor `Sync`. +pub struct EmbeddingEngine { + _inner: OwnedEngine, } impl std::fmt::Debug for Engine { @@ -32,9 +70,14 @@ impl std::fmt::Debug for Engine { } } -/// Builder for one complete serving engine. +/// Builder for one complete text-generation engine. #[derive(Clone, Debug)] pub struct EngineBuilder { + config: ModelConfig, +} + +#[derive(Clone, Debug)] +struct ModelConfig { model_path: PathBuf, tokenizer_config_path: Option, block_size: Option, @@ -46,9 +89,180 @@ pub struct EngineBuilder { speculative_config: Option, prefix_caching: Toggle, max_num_batched_tokens: Option, - scheduler: SchedulerPolicy, + scheduler: Option, kv_transfer_config: Option, jump_forward: Toggle, + device: Option, + gpu_memory_utilization: Option, + kv_cache_memory_bytes: Option, +} + +struct MarshaledModelParams { + raw: Option, + model_path: CString, + tokenizer_config_path: Option, + block_size: Option, + num_blocks: Option, + max_model_len: Option, + max_num_seqs: Option, + tool_parser: Option, + reasoning_parser: Option, + speculative_config: Option, + prefix_caching: Toggle, + max_num_batched_tokens: Option, + scheduling_policy: Option, + kv_transfer_config: Option, + jump_forward: Toggle, + device: Option, + gpu_memory_utilization: Option, + kv_cache_memory_bytes: Option, +} + +impl ModelConfig { + fn new(model_path: impl Into) -> Self { + Self { + model_path: model_path.into(), + tokenizer_config_path: None, + block_size: None, + num_blocks: None, + max_model_len: None, + max_num_seqs: None, + tool_parser: None, + reasoning_parser: None, + speculative_config: None, + prefix_caching: Toggle::Default, + max_num_batched_tokens: None, + scheduler: None, + kv_transfer_config: None, + jump_forward: Toggle::Default, + device: None, + gpu_memory_utilization: None, + kv_cache_memory_bytes: None, + } + } +} + +impl MarshaledModelParams { + fn new(config: ModelConfig) -> Result { + let gpu_memory_utilization = match config.gpu_memory_utilization { + Some(value) if !value.is_finite() || value <= 0.0 => { + return Err(invalid_configuration( + "gpu_memory_utilization must be finite and strictly positive", + )); + } + value => value, + }; + let kv_cache_memory_bytes = match config.kv_cache_memory_bytes { + Some(0) => { + return Err(invalid_configuration( + "kv_cache_memory_bytes must be greater than zero", + )); + } + Some(value) => Some(i64::try_from(value).map_err(|_| { + invalid_configuration("kv_cache_memory_bytes exceeds native i64 range") + })?), + None => None, + }; + + Ok(Self { + raw: None, + model_path: path_to_cstring(&config.model_path, "model path")?, + tokenizer_config_path: config + .tokenizer_config_path + .as_deref() + .map(|path| path_to_cstring(path, "tokenizer config path")) + .transpose()?, + block_size: optional_u32_to_i32(config.block_size, "block_size")?, + num_blocks: optional_u32_to_i32(config.num_blocks, "num_blocks")?, + max_model_len: optional_u32_to_i32(config.max_model_len, "max_model_len")?, + max_num_seqs: optional_u32_to_i32(config.max_num_seqs, "max_num_seqs")?, + tool_parser: optional_cstring(config.tool_parser.as_deref(), "tool parser")?, + reasoning_parser: optional_cstring( + config.reasoning_parser.as_deref(), + "reasoning parser", + )?, + speculative_config: optional_cstring( + config.speculative_config.as_deref(), + "speculative configuration", + )?, + prefix_caching: config.prefix_caching, + max_num_batched_tokens: optional_u32_to_i32( + config.max_num_batched_tokens, + "max_num_batched_tokens", + )?, + scheduling_policy: config + .scheduler + .map(|value| to_cstring(value.as_str(), "scheduler policy")) + .transpose()?, + kv_transfer_config: optional_cstring( + config.kv_transfer_config.as_deref(), + "KV transfer configuration", + )?, + jump_forward: config.jump_forward, + device: config.device, + gpu_memory_utilization, + kv_cache_memory_bytes, + }) + } + + fn apply_defaults(&mut self, mut raw: ffi::vllm_model_params) { + raw.model_path = self.model_path.as_ptr(); + if let Some(value) = &self.tokenizer_config_path { + raw.tokenizer_config_path = value.as_ptr(); + } + if let Some(value) = self.block_size { + raw.block_size = value; + } + if let Some(value) = self.num_blocks { + raw.num_blocks = value; + } + if let Some(value) = self.max_model_len { + raw.max_model_len = value; + } + if let Some(value) = self.max_num_seqs { + raw.max_num_seqs = value; + } + if let Some(value) = &self.tool_parser { + raw.tool_parser = value.as_ptr(); + } + if let Some(value) = &self.reasoning_parser { + raw.reasoning_parser = value.as_ptr(); + } + if let Some(value) = &self.speculative_config { + raw.speculative_config = value.as_ptr(); + } + if self.prefix_caching != Toggle::Default { + raw.enable_prefix_caching = self.prefix_caching.as_native(); + } + if let Some(value) = self.max_num_batched_tokens { + raw.max_num_batched_tokens = value; + } + if let Some(value) = &self.scheduling_policy { + raw.scheduling_policy = value.as_ptr(); + } + if let Some(value) = &self.kv_transfer_config { + raw.kv_transfer_config = value.as_ptr(); + } + if self.jump_forward != Toggle::Default { + raw.enable_jump_forward = self.jump_forward.as_native(); + } + if let Some(value) = self.device { + raw.device = value.as_native(); + } + if let Some(value) = self.gpu_memory_utilization { + raw.gpu_memory_utilization = value; + } + if let Some(value) = self.kv_cache_memory_bytes { + raw.kv_cache_memory_bytes = value; + } + self.raw = Some(raw); + } + + fn raw(&self) -> &ffi::vllm_model_params { + self.raw + .as_ref() + .expect("native defaults must be applied before model loading") + } } /// Why native generation finished. @@ -87,7 +301,7 @@ impl Engine { /// Runs one blocking text completion. pub fn complete(&self, prompt: &str, params: &SamplingParams) -> Result { let prompt = to_cstring(prompt, "prompt")?; - let params = params.marshal()?; + let params = params.marshal(&self.inner.compatibility)?; let mut raw = MaybeUninit::::uninit(); // SAFETY: the engine is owned and live, all pointers remain valid for the // call, and out storage is initialized by native code on success. @@ -129,7 +343,7 @@ impl Engine { F: FnMut(StreamEvent) -> StreamControl, { let prompt = to_cstring(prompt, "prompt")?; - let params = params.marshal()?; + let params = params.marshal(&self.inner.compatibility)?; let mut state = CallbackState::new(&mut callback); // SAFETY: state has a stable stack address for this blocking call; the C // API does not retain user_data after returning. @@ -220,99 +434,97 @@ impl Engine { } } -impl Drop for EngineInner { +impl Drop for OwnedEngine { fn drop(&mut self) { - // SAFETY: EngineInner exclusively owns this live handle. Native teardown - // joins engine workers before returning. + // SAFETY: each OwnedEngine exclusively owns one live handle. Native + // teardown joins engine workers before returning. unsafe { ffi::vllm_engine_free(self.raw.as_ptr()) }; } } -// SAFETY: vllm.cpp documents concurrent completion submissions as thread-safe, -// and EngineInner keeps the engine alive until the last shared owner is dropped. -unsafe impl Send for EngineInner {} -// SAFETY: shared references may submit concurrently through native AsyncLLM; -// destruction cannot race because Arc retains the handle for each active owner. -unsafe impl Sync for EngineInner {} +impl From> for OwnedEngine { + fn from(loaded: LoadedEngine) -> Self { + Self { + raw: loaded.raw, + compatibility: loaded.compatibility, + _task: PhantomData, + _not_send_sync: PhantomData, + } + } +} + +// SAFETY: vllm.cpp documents concurrent text completion submissions as +// thread-safe, and Arc keeps the text engine live through active operations. +unsafe impl Send for OwnedEngine {} +// SAFETY: shared text owners submit through native AsyncLLM; Arc prevents +// destruction from racing an operation. No other task owner receives this impl. +unsafe impl Sync for OwnedEngine {} impl EngineBuilder { #[must_use] pub fn new(model_path: impl Into) -> Self { Self { - model_path: model_path.into(), - tokenizer_config_path: None, - block_size: None, - num_blocks: None, - max_model_len: None, - max_num_seqs: None, - tool_parser: None, - reasoning_parser: None, - speculative_config: None, - prefix_caching: Toggle::Default, - max_num_batched_tokens: None, - scheduler: SchedulerPolicy::Fcfs, - kv_transfer_config: None, - jump_forward: Toggle::Default, + config: ModelConfig::new(model_path), } } #[must_use] pub fn tokenizer_config_path(mut self, value: impl Into) -> Self { - self.tokenizer_config_path = Some(value.into()); + self.config.tokenizer_config_path = Some(value.into()); self } #[must_use] pub fn block_size(mut self, value: u32) -> Self { - self.block_size = Some(value); + self.config.block_size = Some(value); self } #[must_use] pub fn num_blocks(mut self, value: u32) -> Self { - self.num_blocks = Some(value); + self.config.num_blocks = Some(value); self } #[must_use] pub fn max_model_len(mut self, value: u32) -> Self { - self.max_model_len = Some(value); + self.config.max_model_len = Some(value); self } #[must_use] pub fn max_num_seqs(mut self, value: u32) -> Self { - self.max_num_seqs = Some(value); + self.config.max_num_seqs = Some(value); self } #[must_use] pub fn tool_parser(mut self, value: impl Into) -> Self { - self.tool_parser = Some(value.into()); + self.config.tool_parser = Some(value.into()); self } #[must_use] pub fn reasoning_parser(mut self, value: impl Into) -> Self { - self.reasoning_parser = Some(value.into()); + self.config.reasoning_parser = Some(value.into()); self } #[must_use] pub fn speculative_config(mut self, value: impl Into) -> Self { - self.speculative_config = Some(value.into()); + self.config.speculative_config = Some(value.into()); self } #[must_use] pub fn prefix_caching(mut self, value: Toggle) -> Self { - self.prefix_caching = value; + self.config.prefix_caching = value; self } #[must_use] pub fn max_num_batched_tokens(mut self, value: u32) -> Self { - self.max_num_batched_tokens = Some(value); + self.config.max_num_batched_tokens = Some(value); self } @@ -325,97 +537,123 @@ impl EngineBuilder { /// future C ABI/API change. #[must_use] pub fn scheduler(mut self, value: SchedulerPolicy) -> Self { - self.scheduler = value; + self.config.scheduler = Some(value); self } #[must_use] pub fn kv_transfer_config(mut self, value: impl Into) -> Self { - self.kv_transfer_config = Some(value.into()); + self.config.kv_transfer_config = Some(value.into()); self } #[must_use] pub fn jump_forward(mut self, value: Toggle) -> Self { - self.jump_forward = value; + self.config.jump_forward = value; + self + } + + /// Selects the required native device. + /// + /// [`Device::Cuda`] never silently falls back when CUDA is unavailable. + #[must_use] + pub fn device(mut self, value: Device) -> Self { + self.config.device = Some(value); + self + } + + /// Sets the native fraction used by GPU memory profiling. + /// + /// The value must be finite and strictly positive. Values above `1.0` are + /// forwarded because the native ABI does not impose an upper bound. An + /// explicit block count takes precedence over absolute KV-cache bytes, which + /// take precedence over this utilization/profile setting. + #[must_use] + pub fn gpu_memory_utilization(mut self, value: f64) -> Self { + self.config.gpu_memory_utilization = Some(value); + self + } + + /// Sets an absolute KV-cache memory budget in bytes. + /// + /// The value must be nonzero and fit the native signed 64-bit field. Native + /// code validates the model-dependent minimum. An explicit block count takes + /// precedence over this budget, and this budget takes precedence over GPU + /// utilization/profile sizing. + #[must_use] + pub fn kv_cache_memory_bytes(mut self, value: u64) -> Self { + self.config.kv_cache_memory_bytes = Some(value); self } pub fn load(self) -> Result { - let model_path = path_to_cstring(&self.model_path, "model path")?; - let tokenizer_config_path = self - .tokenizer_config_path - .as_deref() - .map(|path| path_to_cstring(path, "tokenizer config path")) - .transpose()?; - let tool_parser = optional_cstring(self.tool_parser.as_deref(), "tool parser")?; - let reasoning_parser = - optional_cstring(self.reasoning_parser.as_deref(), "reasoning parser")?; - let speculative_config = optional_cstring( - self.speculative_config.as_deref(), - "speculative configuration", - )?; - let scheduling_policy = to_cstring(self.scheduler.as_str(), "scheduler policy")?; - let kv_transfer_config = optional_cstring( - self.kv_transfer_config.as_deref(), - "KV transfer configuration", - )?; - - let mut raw = checked_model_params_default()?; - raw.model_path = model_path.as_ptr(); - raw.tokenizer_config_path = optional_pointer(tokenizer_config_path.as_ref()); - raw.block_size = optional_u32_to_i32(self.block_size, "block_size")?; - raw.num_blocks = optional_u32_to_i32(self.num_blocks, "num_blocks")?; - raw.max_model_len = optional_u32_to_i32(self.max_model_len, "max_model_len")?; - raw.max_num_seqs = optional_u32_to_i32(self.max_num_seqs, "max_num_seqs")?; - raw.tool_parser = optional_pointer(tool_parser.as_ref()); - raw.reasoning_parser = optional_pointer(reasoning_parser.as_ref()); - raw.speculative_config = optional_pointer(speculative_config.as_ref()); - raw.enable_prefix_caching = self.prefix_caching.as_native(); - raw.max_num_batched_tokens = - optional_u32_to_i32(self.max_num_batched_tokens, "max_num_batched_tokens")?; - raw.scheduling_policy = scheduling_policy.as_ptr(); - raw.kv_transfer_config = optional_pointer(kv_transfer_config.as_ref()); - raw.enable_jump_forward = self.jump_forward.as_native(); - - let mut output = ptr::null_mut(); - // SAFETY: all string storage remains live for the call and output points - // to writable handle storage. - let status = unsafe { ffi::vllm_engine_load(&raw, &mut output) }; - status_result(status)?; - let raw = NonNull::new(output).ok_or_else(|| Error::ModelLoad { - message: "vllm_engine_load succeeded without a handle".to_owned(), - })?; Ok(Engine { - inner: Arc::new(EngineInner { raw }), + inner: Arc::new(load_engine::(self.config)?), }) } } -fn checked_model_params_default() -> Result { - checked_model_params_default_with( - || { - // SAFETY: this base ABI function takes no pointers or versioned structs. - unsafe { ffi::vllm_abi_version() } - }, - || { - // SAFETY: exact ABI equality was established immediately before this - // by-value return of a versioned struct. - unsafe { ffi::vllm_model_params_default() } +impl TranscriptionEngine { + /// Loads a native engine owner with a transcription-only Rust method surface. + /// + /// ABI 17 cannot inspect the resolved task at load time. This constructor does + /// not probe or infer checkpoint architecture; native task selection remains + /// authoritative for future operations and diagnostics. + pub fn load(model_path: impl Into) -> Result { + Ok(Self { + _inner: load_engine::(ModelConfig::new(model_path))?, + }) + } +} + +impl EmbeddingEngine { + /// Loads a native engine owner with an embedding-only Rust method surface. + /// + /// ABI 17 cannot inspect the resolved task at load time. This constructor does + /// not probe or infer checkpoint architecture; native task selection remains + /// authoritative for future operations and diagnostics. + pub fn load(model_path: impl Into) -> Result { + Ok(Self { + _inner: load_engine::(ModelConfig::new(model_path))?, + }) + } +} + +fn load_engine(config: ModelConfig) -> Result, Error> { + load_engine_with( + config, + Compatibility::check, + |compatibility| compatibility.model_params_default(), + |params, output| { + // SAFETY: the marshaled storage backing every pointer remains live for + // the call and output points to writable handle storage. + unsafe { ffi::vllm_engine_load(params, output) } }, ) + .map(OwnedEngine::from) } -fn checked_model_params_default_with( - abi_version: impl FnOnce() -> i32, - model_params_default: impl FnOnce() -> ffi::vllm_model_params, -) -> Result { - let actual = abi_version(); - let expected = ffi::VLLM_ABI_VERSION as i32; - if actual != expected { - return Err(Error::AbiMismatch { expected, actual }); - } - Ok(model_params_default()) +fn load_engine_with( + config: ModelConfig, + check: impl FnOnce() -> Result, + defaults: impl FnOnce(&Compatibility) -> ffi::vllm_model_params, + load: impl FnOnce(&ffi::vllm_model_params, *mut *mut ffi::vllm_engine) -> ffi::vllm_status, +) -> Result, Error> { + let mut params = MarshaledModelParams::new(config)?; + let compatibility = check()?; + params.apply_defaults(defaults(&compatibility)); + + let mut output = ptr::null_mut(); + let status = load(params.raw(), &mut output); + status_result(status)?; + let raw = NonNull::new(output).ok_or_else(|| Error::ModelLoad { + message: "vllm_engine_load succeeded without a handle".to_owned(), + })?; + Ok(LoadedEngine { + raw, + compatibility, + _task: PhantomData, + }) } fn completion_from_raw(raw: &ffi::vllm_completion) -> Result { @@ -457,22 +695,19 @@ fn count_to_u32(value: i32, field: &'static str) -> Result { u32::try_from(value).map_err(|_| invalid_configuration(format!("native {field} was negative"))) } -fn optional_u32_to_i32(value: Option, field: &'static str) -> Result { - match value { - Some(0) | None => Ok(0), - Some(value) => i32::try_from(value) - .map_err(|_| invalid_configuration(format!("{field} exceeds native i32 range"))), - } +fn optional_u32_to_i32(value: Option, field: &'static str) -> Result, Error> { + value + .map(|value| { + i32::try_from(value) + .map_err(|_| invalid_configuration(format!("{field} exceeds native i32 range"))) + }) + .transpose() } fn optional_cstring(value: Option<&str>, field: &'static str) -> Result, Error> { value.map(|value| to_cstring(value, field)).transpose() } -fn optional_pointer(value: Option<&CString>) -> *const c_char { - value.map_or(ptr::null(), |value| value.as_ptr()) -} - fn to_cstring(value: &str, field: &'static str) -> Result { CString::new(value).map_err(|_| Error::InteriorNul { field }) } @@ -523,31 +758,214 @@ impl Drop for NativeStringGuard { #[cfg(test)] mod tests { - use super::checked_model_params_default_with; - use crate::Error; use std::cell::RefCell; + use std::ffi::CStr; + use std::ptr::{self, NonNull}; + + use vllm_cpp_sys as ffi; + + use super::{ + load_engine_with, Device, EmbeddingTask, MarshaledModelParams, ModelConfig, + SchedulerPolicy, TextTask, Toggle, TranscriptionTask, + }; + use crate::abi::Compatibility; + use crate::Error; + + const NATIVE_STRING: &[u8] = b"native-default\0"; + + fn native_defaults() -> ffi::vllm_model_params { + let pointer = NATIVE_STRING.as_ptr().cast(); + ffi::vllm_model_params { + model_path: pointer, + tokenizer_config_path: pointer, + block_size: 41, + num_blocks: 42, + max_model_len: 43, + max_num_seqs: 44, + tool_parser: pointer, + reasoning_parser: pointer, + speculative_config: pointer, + enable_prefix_caching: 45, + max_num_batched_tokens: 46, + scheduling_policy: pointer, + kv_transfer_config: pointer, + enable_jump_forward: 47, + device: 48, + gpu_memory_utilization: 0.92, + kv_cache_memory_bytes: 49, + } + } + + fn matching_compatibility() -> Result { + Compatibility::check_with(|| ffi::VLLM_ABI_VERSION as i32) + } + + fn c_string(pointer: *const std::os::raw::c_char) -> String { + assert!(!pointer.is_null()); + // SAFETY: tests inspect pointers while their MarshaledModelParams owner is live. + unsafe { CStr::from_ptr(pointer) } + .to_str() + .expect("UTF-8 test string") + .to_owned() + } #[test] - fn abi_mismatch_prevents_model_params_default_call() { + fn default_application_preserves_every_unset_native_value() { + let defaults = native_defaults(); + let mut params = MarshaledModelParams::new(ModelConfig::new("model-dir")) + .expect("marshal default model config"); + params.apply_defaults(defaults); + let raw = params.raw(); + + assert_eq!(c_string(raw.model_path), "model-dir"); + assert_eq!(raw.tokenizer_config_path, defaults.tokenizer_config_path); + assert_eq!(raw.block_size, defaults.block_size); + assert_eq!(raw.num_blocks, defaults.num_blocks); + assert_eq!(raw.max_model_len, defaults.max_model_len); + assert_eq!(raw.max_num_seqs, defaults.max_num_seqs); + assert_eq!(raw.tool_parser, defaults.tool_parser); + assert_eq!(raw.reasoning_parser, defaults.reasoning_parser); + assert_eq!(raw.speculative_config, defaults.speculative_config); + assert_eq!(raw.enable_prefix_caching, defaults.enable_prefix_caching); + assert_eq!(raw.max_num_batched_tokens, defaults.max_num_batched_tokens); + assert_eq!(raw.scheduling_policy, defaults.scheduling_policy); + assert_eq!(raw.kv_transfer_config, defaults.kv_transfer_config); + assert_eq!(raw.enable_jump_forward, defaults.enable_jump_forward); + assert_eq!(raw.device, defaults.device); + assert_eq!(raw.gpu_memory_utilization, defaults.gpu_memory_utilization); + assert_eq!(raw.kv_cache_memory_bytes, defaults.kv_cache_memory_bytes); + } + + #[test] + fn explicit_overrides_and_strings_reach_the_native_view() { + let mut config = ModelConfig::new("model-dir"); + config.tokenizer_config_path = Some("tokenizer.json".into()); + config.block_size = Some(16); + config.num_blocks = Some(32); + config.max_model_len = Some(128); + config.max_num_seqs = Some(2); + config.tool_parser = Some("hermes".to_owned()); + config.reasoning_parser = Some("reasoning".to_owned()); + config.speculative_config = Some("{}".to_owned()); + config.prefix_caching = Toggle::Off; + config.max_num_batched_tokens = Some(64); + config.scheduler = Some(SchedulerPolicy::LongestPrefixMatch); + config.kv_transfer_config = Some("{\"kv_role\":\"kv_both\"}".to_owned()); + config.jump_forward = Toggle::On; + config.device = Some(Device::Cuda); + config.gpu_memory_utilization = Some(1.25); + config.kv_cache_memory_bytes = Some(4096); + + let mut params = MarshaledModelParams::new(config).expect("marshal explicit model config"); + params.apply_defaults(native_defaults()); + let raw = params.raw(); + + assert_eq!(c_string(raw.model_path), "model-dir"); + assert_eq!(c_string(raw.tokenizer_config_path), "tokenizer.json"); + assert_eq!(raw.block_size, 16); + assert_eq!(raw.num_blocks, 32); + assert_eq!(raw.max_model_len, 128); + assert_eq!(raw.max_num_seqs, 2); + assert_eq!(c_string(raw.tool_parser), "hermes"); + assert_eq!(c_string(raw.reasoning_parser), "reasoning"); + assert_eq!(c_string(raw.speculative_config), "{}"); + assert_eq!(raw.enable_prefix_caching, Toggle::Off.as_native()); + assert_eq!(raw.max_num_batched_tokens, 64); + assert_eq!(c_string(raw.scheduling_policy), "lpm"); + assert_eq!( + c_string(raw.kv_transfer_config), + "{\"kv_role\":\"kv_both\"}" + ); + assert_eq!(raw.enable_jump_forward, Toggle::On.as_native()); + assert_eq!(raw.device, Device::Cuda.as_native()); + assert_eq!(raw.gpu_memory_utilization, 1.25); + assert_eq!(raw.kv_cache_memory_bytes, 4096); + } + + #[test] + fn forwards_all_memory_settings_without_changing_native_precedence() { + let mut config = ModelConfig::new("model-dir"); + config.num_blocks = Some(7); + config.gpu_memory_utilization = Some(2.0); + config.kv_cache_memory_bytes = Some(8192); + + let mut params = MarshaledModelParams::new(config).expect("marshal memory settings"); + params.apply_defaults(native_defaults()); + let raw = params.raw(); + assert_eq!(raw.num_blocks, 7); + assert_eq!(raw.kv_cache_memory_bytes, 8192); + assert_eq!(raw.gpu_memory_utilization, 2.0); + } + + fn assert_shared_load_order() { let calls = RefCell::new(Vec::new()); - let result = checked_model_params_default_with( + let loaded = load_engine_with::( + ModelConfig::new("shared-model"), || { calls.borrow_mut().push("abi"); - 10 + matching_compatibility() }, - || { + |_| { calls.borrow_mut().push("default"); - unreachable!("default helper must not run after an ABI mismatch") + native_defaults() + }, + |params, output| { + calls.borrow_mut().push("load"); + assert_eq!(c_string(params.model_path), "shared-model"); + // SAFETY: output is writable storage supplied by load_engine_with; + // the dangling non-null value is never dereferenced or freed. + unsafe { *output = NonNull::::dangling().as_ptr() }; + ffi::vllm_status_VLLM_OK }, + ) + .expect("injected load"); + + assert_eq!(*calls.borrow(), ["abi", "default", "load"]); + assert_eq!(loaded.raw, NonNull::dangling()); + } + + #[test] + fn shared_loader_checks_abi_before_defaults_for_every_task_marker() { + assert_shared_load_order::(); + assert_shared_load_order::(); + assert_shared_load_order::(); + } + + #[test] + fn rust_marshaling_failure_precedes_the_abi_probe() { + let calls = RefCell::new(Vec::new()); + let result = load_engine_with::( + ModelConfig::new("bad\0model"), + || { + calls.borrow_mut().push("abi"); + matching_compatibility() + }, + |_| unreachable!("default helper must not run after marshaling failure"), + |_, _| unreachable!("load must not run after marshaling failure"), ); assert!(matches!( result, - Err(Error::AbiMismatch { - expected: 17, - actual: 10 + Err(Error::InteriorNul { + field: "model path" }) )); - assert_eq!(*calls.borrow(), ["abi"]); + assert!(calls.borrow().is_empty()); + } + + #[test] + fn successful_status_with_null_handle_is_rejected() { + let result = load_engine_with::( + ModelConfig::new("model-dir"), + matching_compatibility, + |_| native_defaults(), + |_, output| { + // SAFETY: output is writable storage supplied by load_engine_with. + unsafe { *output = ptr::null_mut() }; + ffi::vllm_status_VLLM_OK + }, + ); + + assert!(matches!(result, Err(Error::ModelLoad { .. }))); } } diff --git a/vllm-cpp/src/lib.rs b/vllm-cpp/src/lib.rs index 794782c..952d81b 100644 --- a/vllm-cpp/src/lib.rs +++ b/vllm-cpp/src/lib.rs @@ -3,12 +3,14 @@ //! # Entry points //! //! Resolve a Hub model with [`HuggingFaceModel`] (default `main`, or an explicit -//! revision), then create an [`Engine`] with [`Engine::load`] or configure native -//! model settings through [`EngineBuilder`]. [`SamplingParams`] owns sampling, +//! revision), then create a text [`Engine`] with [`Engine::load`] or configure +//! native model settings through [`EngineBuilder`]. Path-only +//! [`TranscriptionEngine`] and [`EmbeddingEngine`] owners reserve task-specific +//! method surfaces for future operations. [`SamplingParams`] owns sampling, //! stop-string, [`StructuredOutput`] settings, and an optional host-side logits -//! processor for completion calls. The engine provides -//! blocking completion, streaming, raw-JSON chat, and [`Engine::submit`] for a -//! concurrent [`Request`]. Enable `serde` for `serde_json::Value` chat helpers. +//! processor for completion calls. The text engine provides blocking completion, +//! streaming, raw-JSON chat, and [`Engine::submit`] for a concurrent [`Request`]. +//! Enable `serde` for `serde_json::Value` chat helpers. //! //! # Ownership and callbacks //! @@ -21,12 +23,15 @@ //! [`Error::LogitsProcessorPanicked`]. //! //! A [`Request`] retains its engine and asynchronous callback until native -//! free/join completes. Requests are `Send` but intentionally not `Sync`, while -//! engines are `Send + Sync`. Asynchronous callbacks run on a native delivery -//! thread, must be `Send + 'static`, and surface panic through +//! free/join completes. Requests are `Send` but intentionally not `Sync`. The text +//! [`Engine`] is `Send + Sync`; transcription and embedding owners are +//! conservatively neither and must remain on their creating thread. Asynchronous +//! callbacks run on a native delivery thread, must be `Send + 'static`, and surface panic through //! [`Error::CallbackPanicked`]. ABI version 17 forbids waiting for or freeing a //! request from its callback thread; callback-thread drop delegates ownership to -//! a cleanup reaper instead. +//! a cleanup reaper instead. ABI 17 exposes no task-introspection API, so loading +//! cannot prove or infer a checkpoint's task. Native task selection and future +//! wrong-task diagnostics remain authoritative. //! //! # ABI, linking, and deployment //! @@ -47,6 +52,7 @@ //! Accelerator features are build/configuration surfaces with known runtime //! blockers, not complete accelerator runtime support. +mod abi; mod callback; mod engine; mod error; @@ -55,10 +61,12 @@ mod params; mod request; pub use callback::{StreamControl, StreamEvent, StreamOutcome}; -pub use engine::{Completion, Engine, EngineBuilder, FinishReason}; +pub use engine::{ + Completion, EmbeddingEngine, Engine, EngineBuilder, FinishReason, TranscriptionEngine, +}; pub use error::{Error, HuggingFaceError}; pub use hf::HuggingFaceModel; -pub use params::{SamplingParams, SchedulerPolicy, StructuredOutput, Toggle}; +pub use params::{Device, SamplingParams, SchedulerPolicy, StructuredOutput, Toggle}; pub use request::{Request, RequestOutcome}; /// Returns the compile-time C ABI expected by this crate. diff --git a/vllm-cpp/src/params.rs b/vllm-cpp/src/params.rs index 863e1a9..c094b83 100644 --- a/vllm-cpp/src/params.rs +++ b/vllm-cpp/src/params.rs @@ -12,10 +12,37 @@ use std::thread::{self, ThreadId}; use vllm_cpp_sys as ffi; +use crate::abi::Compatibility; use crate::error::{invalid_configuration, Error}; const NATIVE_DEFAULT_MAX_TOKENS: u32 = 16; +/// Device selection for a text, transcription, or embedding engine. +/// +/// [`Auto`](Self::Auto) preserves native platform selection. [`Cuda`](Self::Cuda) +/// requires the CUDA platform; native loading fails rather than falling back to +/// another device when CUDA is unavailable. +#[derive(Clone, Copy, Debug, Default, Eq, PartialEq)] +pub enum Device { + /// Let vllm.cpp select the available platform. + #[default] + Auto, + /// Require the CPU platform. + Cpu, + /// Require the CUDA platform without fallback. + Cuda, +} + +impl Device { + pub(crate) const fn as_native(self) -> i32 { + match self { + Self::Auto => 0, + Self::Cpu => 1, + Self::Cuda => 2, + } + } +} + /// Native scheduler admission order. /// /// Raw and serde chat request JSON can carry a `priority` field that the native @@ -278,8 +305,11 @@ impl SamplingParams { self } - pub(crate) fn marshal(&self) -> Result { - MarshaledSamplingParams::new(self) + pub(crate) fn marshal( + &self, + compatibility: &Compatibility, + ) -> Result { + MarshaledSamplingParams::new(self, compatibility) } } @@ -294,9 +324,8 @@ pub(crate) struct MarshaledSamplingParams { } impl MarshaledSamplingParams { - fn new(params: &SamplingParams) -> Result { - // ABI equality is checked before this struct-returning call. - let mut raw = unsafe { ffi::vllm_sampling_params_default() }; + fn new(params: &SamplingParams, compatibility: &Compatibility) -> Result { + let mut raw = compatibility.sampling_params_default(); raw.temperature = params.temperature; raw.top_p = params.top_p; raw.top_k = params.top_k; @@ -611,11 +640,24 @@ fn length_to_i32(value: usize, field: &'static str) -> Result { #[cfg(test)] mod tests { - use super::{logits_processor_trampoline, SamplingParams}; + use super::{logits_processor_trampoline, Device, SamplingParams}; + use crate::abi::Compatibility; use crate::Error; use std::sync::atomic::{AtomicUsize, Ordering}; use std::sync::Arc; + fn compatibility() -> Compatibility { + Compatibility::check().expect("matching native ABI") + } + + #[test] + fn maps_devices_to_native_values() { + assert_eq!(Device::default(), Device::Auto); + assert_eq!(Device::Auto.as_native(), 0); + assert_eq!(Device::Cpu.as_native(), 1); + assert_eq!(Device::Cuda.as_native(), 2); + } + #[test] fn marshals_and_invokes_custom_logits_processor() { let params = SamplingParams::default() @@ -624,7 +666,7 @@ mod tests { assert_eq!(tokens, &[3, 5]); logits[1] = 9.0; }); - let marshaled = params.marshal().expect("marshal processor"); + let marshaled = params.marshal(&compatibility()).expect("marshal processor"); let callback = marshaled .raw() .logits_processor @@ -654,7 +696,7 @@ mod tests { .logits_processor(move |_, _| { calls.fetch_add(1, Ordering::Relaxed); }); - let marshaled = params.marshal().expect("marshal processor"); + let marshaled = params.marshal(&compatibility()).expect("marshal processor"); marshaled.raw().logits_processor_user_data }; let mut logits = [1.0]; @@ -675,7 +717,7 @@ mod tests { let params = SamplingParams::default() .max_tokens(2) .logits_processor(|_, _| panic!("processor panic")); - let marshaled = params.marshal().expect("marshal processor"); + let marshaled = params.marshal(&compatibility()).expect("marshal processor"); let mut logits = [1.0, 2.0]; unsafe { logits_processor_trampoline( @@ -709,12 +751,15 @@ mod tests { .unbounded() .logits_processor(|_, _| {}), ] { - let error = params.marshal().err().expect("invalid bounds rejection"); + let error = params + .marshal(&compatibility()) + .err() + .expect("invalid bounds rejection"); assert!(matches!(error, Error::InvalidConfiguration { .. })); } let params = SamplingParams::default().logits_processor(|_, _| {}); - let marshaled = params.marshal().expect("marshal processor"); + let marshaled = params.marshal(&compatibility()).expect("marshal processor"); unsafe { logits_processor_trampoline( std::ptr::null(), diff --git a/vllm-cpp/src/request.rs b/vllm-cpp/src/request.rs index 4c35c46..91b02f1 100644 --- a/vllm-cpp/src/request.rs +++ b/vllm-cpp/src/request.rs @@ -73,7 +73,7 @@ impl Engine { { cleanup_sender()?; let prompt = to_cstring(prompt, "prompt")?; - let mut params = params.marshal()?; + let mut params = params.marshal(&self.inner.compatibility)?; let mut callback = Box::new(AsyncCallbackState::new(callback)); let mut output = ptr::null_mut(); // SAFETY: the engine is retained by the returned Request, native code diff --git a/vllm-cpp/tests/safe_api.rs b/vllm-cpp/tests/safe_api.rs index 13c0e06..94ced05 100644 --- a/vllm-cpp/tests/safe_api.rs +++ b/vllm-cpp/tests/safe_api.rs @@ -1,14 +1,19 @@ use static_assertions::{assert_impl_all, assert_not_impl_any}; use vllm_cpp::{ - Engine, Error, HuggingFaceError, HuggingFaceModel, Request, SchedulerPolicy, Toggle, + Device, EmbeddingEngine, Engine, EngineBuilder, Error, HuggingFaceError, HuggingFaceModel, + Request, SchedulerPolicy, Toggle, TranscriptionEngine, }; +assert_impl_all!(Device: Clone, Copy, std::fmt::Debug, Default, Eq, PartialEq, Send, Sync); assert_impl_all!(Engine: Send, Sync, Clone); +assert_impl_all!(EngineBuilder: Clone, std::fmt::Debug, Send, Sync); assert_impl_all!(HuggingFaceError: Clone, std::fmt::Debug, Eq, PartialEq); assert_impl_all!(HuggingFaceModel: Clone, std::fmt::Debug); assert_impl_all!(Request: Send); assert_impl_all!(vllm_cpp::SamplingParams: Clone, Send, Sync); assert_not_impl_any!(Request: Sync); +assert_not_impl_any!(TranscriptionEngine: Send, Sync); +assert_not_impl_any!(EmbeddingEngine: Send, Sync); fn missing_model() -> &'static str { "/nonexistent/vllm-cpp-rs-safe-api-model" @@ -32,10 +37,20 @@ fn reports_expected_abi() { } #[test] -fn missing_model_is_typed() { - let error = Engine::load(missing_model()).unwrap_err(); - assert!(matches!(error, Error::ModelLoad { .. }), "{error:?}"); - assert!(!error.to_string().is_empty()); +fn missing_model_is_typed_for_every_task_owner() { + let errors = [ + Engine::load(missing_model()).unwrap_err(), + TranscriptionEngine::load(missing_model()) + .err() + .expect("missing transcription model error"), + EmbeddingEngine::load(missing_model()) + .err() + .expect("missing embedding model error"), + ]; + for error in errors { + assert!(matches!(error, Error::ModelLoad { .. }), "{error:?}"); + assert!(!error.to_string().is_empty()); + } } #[test] @@ -58,6 +73,11 @@ fn interior_nul_fails_before_ffi() { ); } +#[test] +fn device_defaults_to_native_auto_selection() { + assert_eq!(Device::default(), Device::Auto); +} + #[test] fn engine_builder_accepts_all_safe_options() { let error = Engine::builder(missing_model()) @@ -73,7 +93,65 @@ fn engine_builder_accepts_all_safe_options() { .scheduler(SchedulerPolicy::LongestPrefixMatch) .kv_transfer_config("") .jump_forward(Toggle::Off) + .device(Device::Cpu) + .gpu_memory_utilization(1.25) + .kv_cache_memory_bytes(4096) .load() .unwrap_err(); assert!(matches!(error, Error::ModelLoad { .. }), "{error:?}"); } + +#[test] +fn rejects_invalid_gpu_memory_utilization() { + for value in [f64::NAN, f64::INFINITY, f64::NEG_INFINITY, 0.0, -0.0, -1.0] { + let error = Engine::builder(missing_model()) + .gpu_memory_utilization(value) + .load() + .unwrap_err(); + assert!( + matches!(error, Error::InvalidConfiguration { .. }), + "{value:?}: {error:?}" + ); + } +} + +#[test] +fn rejects_invalid_kv_cache_memory_bytes() { + for value in [0, i64::MAX as u64 + 1, u64::MAX] { + let error = Engine::builder(missing_model()) + .kv_cache_memory_bytes(value) + .load() + .unwrap_err(); + assert!( + matches!(error, Error::InvalidConfiguration { .. }), + "{value}: {error:?}" + ); + } +} + +#[test] +fn valid_memory_settings_and_native_precedence_reach_model_loading() { + let errors = [ + Engine::builder(missing_model()) + .gpu_memory_utilization(f64::MIN_POSITIVE) + .load() + .unwrap_err(), + Engine::builder(missing_model()) + .gpu_memory_utilization(2.0) + .load() + .unwrap_err(), + Engine::builder(missing_model()) + .kv_cache_memory_bytes(i64::MAX as u64) + .load() + .unwrap_err(), + Engine::builder(missing_model()) + .num_blocks(1) + .kv_cache_memory_bytes(4096) + .gpu_memory_utilization(1.5) + .load() + .unwrap_err(), + ]; + for error in errors { + assert!(matches!(error, Error::ModelLoad { .. }), "{error:?}"); + } +}