diff --git a/Cargo.toml b/Cargo.toml index a5c99f1bb..9ae304cfd 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,10 +30,12 @@ axum = { version = "0.8.9" } blake3 = "1.8" lz4_flex = "0.13" bytemuck = { version = "1.25" } +memmap2 = "0.9" chrono = { version = "0.4" } dashmap = "6.2.1" http = "1.4.2" lru = "0.18.0" +libc = "0.2" num-traits = "0.2" parking_lot = "0.12.5" rand = "0.10.1" diff --git a/crates/multimodal/Cargo.toml b/crates/multimodal/Cargo.toml index 7402b7bb5..c1508e992 100644 --- a/crates/multimodal/Cargo.toml +++ b/crates/multimodal/Cargo.toml @@ -28,6 +28,7 @@ image = { version = "0.25.10", default-features = false, features = ["png", "jpe libloading = "0.8" ndarray = "0.17" once_cell = "1.21.4" +rayon = "1.12" opencv = { version = "0.99.0", default-features = false, features = ["clang-runtime", "imgproc", "videoio"], optional = true } reqwest = { workspace = true, features = ["stream"] } serde = { workspace = true, features = ["derive"] } @@ -35,7 +36,15 @@ serde_bytes = "0.11" serde_json.workspace = true tempfile = "3.27" thiserror.workspace = true -tokio = { workspace = true, features = ["sync", "fs", "rt-multi-thread", "process", "time"] } +tokio = { workspace = true, features = [ + "sync", + "fs", + "io-util", + "macros", + "rt-multi-thread", + "process", + "time", +] } tracing.workspace = true url = "2.5.8" @@ -49,6 +58,10 @@ criterion = { version = "0.8", features = ["html_reports"] } npyz = { version = "0.9", features = ["npz"] } tokio = { workspace = true, features = ["rt-multi-thread", "macros"] } +[build-dependencies] +cc = "1" +pkg-config = "0.3" + [[bench]] name = "image_preprocess" harness = false diff --git a/crates/multimodal/build.rs b/crates/multimodal/build.rs new file mode 100644 index 000000000..8ec6387cf --- /dev/null +++ b/crates/multimodal/build.rs @@ -0,0 +1,31 @@ +use std::env; + +fn main() -> Result<(), Box> { + println!("cargo:rerun-if-changed=src/opencv_buffer_capture.cpp"); + println!("cargo:rerun-if-env-changed=OPENCV_INCLUDE_PATHS"); + if env::var_os("CARGO_FEATURE_OPENCV_VIDEO").is_none() { + return Ok(()); + } + + let mut build = cc::Build::new(); + build + .cpp(true) + .file("src/opencv_buffer_capture.cpp") + .flag_if_supported("-std=c++17"); + + if let Some(paths) = env::var_os("OPENCV_INCLUDE_PATHS") { + for path in env::split_paths(&paths) { + build.include(path); + } + } else { + let opencv = pkg_config::Config::new() + .cargo_metadata(false) + .probe("opencv4")?; + for path in opencv.include_paths { + build.include(path); + } + } + + build.compile("smg_opencv_buffer_capture"); + Ok(()) +} diff --git a/crates/multimodal/src/error.rs b/crates/multimodal/src/error.rs index 402b262c2..eb460b574 100644 --- a/crates/multimodal/src/error.rs +++ b/crates/multimodal/src/error.rs @@ -24,6 +24,8 @@ pub enum MediaConnectorError { DataUrl(String), #[error("media decode task failed: {0}")] Blocking(#[from] tokio::task::JoinError), + #[error("failed to initialize multimodal runtime: {0}")] + Runtime(String), #[error("image decode error: {0}")] Image(#[from] image::ImageError), #[error("video decode error: {0}")] diff --git a/crates/multimodal/src/lib.rs b/crates/multimodal/src/lib.rs index 7e63f5558..15e8f8e95 100644 --- a/crates/multimodal/src/lib.rs +++ b/crates/multimodal/src/lib.rs @@ -3,7 +3,10 @@ pub mod hasher; pub mod hub; pub mod jpeg_turbo; pub mod media; +#[cfg(feature = "opencv-video")] +mod opencv_buffer; pub mod registry; +pub mod runtime; pub mod tracker; pub mod types; pub mod vision; @@ -13,6 +16,7 @@ pub use media::{ ImageFetchConfig, MediaConnector, MediaConnectorConfig, MediaSource, VideoFetchConfig, }; pub use registry::{ModelMetadata, ModelProcessorSpec, ModelRegistry}; +pub use runtime::MultimodalRuntime; pub use tracker::{AsyncMultiModalTracker, TrackerOutput}; pub use types::{ FieldLayout, ImageDetail, ImageFrame, ImageSize, ImageSource, MediaContentPart, Modality, @@ -21,6 +25,8 @@ pub use types::{ }; // Re-export vision processing components pub use vision::{ - LlavaNextProcessor, LlavaProcessor, ModelSpecificValue, PreProcessorConfig, - PreprocessedEncoderInputs, TransformError, VisionPreProcessor, VisionProcessorRegistry, + DeferredNormalizedEncoderInput, EncoderInput, LlavaNextProcessor, LlavaProcessor, + ModalityPreProcessor, ModalityProcessorRegistry, ModelSpecificValue, OutputPreference, + PreProcessorConfig, PreprocessRequest, PreprocessedEncoderInputs, TransformError, VideoInput, + VisionInput, VisionPreProcessor, VisionPreprocessRequest, VisionProcessorRegistry, }; diff --git a/crates/multimodal/src/media.rs b/crates/multimodal/src/media.rs index f0fda9326..53cff4152 100644 --- a/crates/multimodal/src/media.rs +++ b/crates/multimodal/src/media.rs @@ -1,3 +1,5 @@ +#[cfg(feature = "opencv-video")] +use std::sync::mpsc::{sync_channel, SyncSender}; use std::{ collections::HashSet, io::Write, @@ -10,9 +12,14 @@ use std::{ use base64::{engine::general_purpose::STANDARD as BASE64_STANDARD, Engine}; use bytes::Bytes; #[cfg(feature = "opencv-video")] -use opencv::{core::Mat, imgproc, prelude::*, videoio}; +use opencv::{ + core::{Mat, Vector}, + imgproc, + prelude::*, + videoio, +}; use reqwest::Client; -use tokio::{fs, process::Command, task, time}; +use tokio::{fs, io::AsyncReadExt, process::Command, task, time}; use tracing::info; use url::Url; @@ -22,9 +29,26 @@ static VIDEO_DECODE_BACKEND: OnceLock> = OnceLock::new(); static LOG_VIDEO_DECODE_TIMING: OnceLock = OnceLock::new(); static VIDEO_PROCESS_TIMEOUT: OnceLock = OnceLock::new(); static VIDEO_MAX_DECODED_BYTES: OnceLock = OnceLock::new(); +#[cfg(feature = "opencv-video")] +const MAX_OPENCV_DECODER_THREADS: usize = 8; +#[cfg(feature = "opencv-video")] +const OPENCV_DECODE_BURST_COALESCE: Duration = Duration::from_millis(5); +#[cfg(feature = "opencv-video")] +const OPENCV_LOW_CONCURRENCY_LIMIT: usize = 8; +#[cfg(feature = "opencv-video")] +const OPENCV_LOW_CONCURRENCY_CPU_MULTIPLIER: usize = 2; +#[cfg(feature = "opencv-video")] +const OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_NUMERATOR: usize = 6; +#[cfg(feature = "opencv-video")] +const OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_DENOMINATOR: usize = 7; +#[cfg(feature = "opencv-video")] +use super::runtime::ActiveVideoDecode; +#[cfg(feature = "opencv-video")] +use super::types::{DecodedRgbFrameStream, OwnedRgbFrame, RgbChannelOrder}; use super::{ error::MediaConnectorError, + runtime::MultimodalRuntime, types::{ DecodedRgbFrame, DecodedRgbVideo, ImageDetail, ImageFrame, ImageSource, VideoClip, VideoSource, @@ -92,10 +116,23 @@ pub struct MediaConnector { allowed_domains: Option>, allowed_local_media_path: Option, fetch_timeout: Duration, + runtime: Arc, } impl MediaConnector { pub fn new(client: Client, config: MediaConnectorConfig) -> Result { + let runtime = Arc::new( + MultimodalRuntime::new() + .map_err(|error| MediaConnectorError::Runtime(error.to_string()))?, + ); + Self::new_with_runtime(client, config, runtime) + } + + pub fn new_with_runtime( + client: Client, + config: MediaConnectorConfig, + runtime: Arc, + ) -> Result { let allowed_domains = config.allowed_domains.map(|domains| { domains .into_iter() @@ -114,6 +151,7 @@ impl MediaConnector { allowed_domains, allowed_local_media_path, fetch_timeout: config.fetch_timeout, + runtime, }) } @@ -201,8 +239,8 @@ impl MediaConnector { } let data = data.trim(); - let decoded = BASE64_STANDARD.decode(data)?; - self.decode_image(decoded.into(), cfg.detail, ImageSource::DataUrl) + let decoded = decode_base64_data_url_payload(data).await?; + self.decode_image(decoded, cfg.detail, ImageSource::DataUrl) .await } @@ -222,9 +260,8 @@ impl MediaConnector { } let data = data.trim(); - let decoded = BASE64_STANDARD.decode(data)?; - self.decode_video(decoded.into(), cfg, VideoSource::DataUrl) - .await + let decoded = decode_base64_data_url_payload(data).await?; + self.decode_video(decoded, cfg, VideoSource::DataUrl).await } async fn fetch_file( @@ -303,9 +340,12 @@ impl MediaConnector { )); } - let bytes = fs::read(&canonical).await?; - self.decode_video(bytes.into(), cfg, VideoSource::File { path: canonical }) - .await + let source = VideoSource::File { + path: canonical.clone(), + }; + + let bytes = Bytes::from(fs::read(&canonical).await?); + self.decode_video(bytes, cfg, source).await } fn ensure_domain_allowed(&self, url: &Url) -> Result<(), MediaConnectorError> { @@ -328,25 +368,7 @@ impl MediaConnector { source: ImageSource, ) -> Result, MediaConnectorError> { let hash = crate::hasher::hash_image(&bytes); - - // Decode JPEGs through libjpeg-turbo (PIL-compatible defaults: accurate - // IDCT + fancy upsampling) so pixel values match vLLM bit-for-bit; the - // pure-Rust decoder diverges by a few levels, which the vision encoder - // amplifies into an embedding shift. Non-JPEG inputs and any turbojpeg - // failure fall back to the `image` crate. - let bytes_for_decode = bytes.clone(); - let image = task::spawn_blocking( - move || -> Result { - if let Some(img) = crate::jpeg_turbo::decode_jpeg_rgb(&bytes_for_decode) { - return Ok(img); - } - let cursor = std::io::Cursor::new(bytes_for_decode); - let reader = image::ImageReader::new(cursor).with_guessed_format()?; - Ok(reader.decode()?) - }, - ) - .await - .map_err(MediaConnectorError::Blocking)??; + let image = decode_image(bytes.clone()).await?; Ok(Arc::new(ImageFrame::new( image, bytes, detail, source, hash, @@ -359,48 +381,168 @@ impl MediaConnector { cfg: VideoFetchConfig, source: VideoSource, ) -> Result, MediaConnectorError> { - if cfg.max_frames == 0 { - return Err(MediaConnectorError::VideoDecode( - "max_frames must be greater than 0".to_string(), - )); - } - if cfg.min_frames == 0 { - return Err(MediaConnectorError::VideoDecode( - "min_frames must be greater than 0".to_string(), - )); - } - if cfg.min_frames > cfg.max_frames { - return Err(MediaConnectorError::VideoDecode( - "min_frames must be less than or equal to max_frames".to_string(), - )); - } - if cfg.sample_fps <= 0.0 { - return Err(MediaConnectorError::VideoDecode( - "sample_fps must be greater than 0".to_string(), - )); - } + validate_video_fetch_config(cfg)?; + let bytes_for_hash = bytes.clone(); + let bytes_for_decode = bytes.clone(); + let hash = async move { + task::spawn_blocking(move || crate::hasher::hash_video(&bytes_for_hash)) + .await + .map_err(MediaConnectorError::Blocking) + }; + let decode = decode_video_frames(bytes_for_decode, cfg, self.runtime.clone()); + let (hash, decoded) = tokio::try_join!(hash, decode)?; - let hash = crate::hasher::hash_video(&bytes); - let decoded = decode_video_frames(bytes.clone(), cfg).await?; + Ok(Arc::new(video_clip_from_decoded( + decoded, bytes, source, hash, + ))) + } +} - let clip = match decoded { - DecodedVideoFrames::Images(frames) => VideoClip::new(frames, bytes, source, hash), - DecodedVideoFrames::Rgb(rgb_video) => { - VideoClip::new_rgb(rgb_video, bytes, source, hash) +async fn decode_base64_data_url_payload(data: &str) -> Result { + let data = data.to_owned(); + let decoded = task::spawn_blocking(move || BASE64_STANDARD.decode(data)) + .await + .map_err(MediaConnectorError::Blocking)??; + Ok(Bytes::from(decoded)) +} + +async fn decode_image(bytes: Bytes) -> Result { + // Decode JPEGs through libjpeg-turbo with PIL-compatible defaults + // (accurate IDCT + fancy upsampling). The pure-Rust decoder can diverge by + // a few pixel levels, which the vision encoder amplifies into an embedding + // shift. Non-JPEG inputs and any turbojpeg failure fall back to the `image` + // crate. + task::spawn_blocking( + move || -> Result { + if let Some(img) = crate::jpeg_turbo::decode_jpeg_rgb(&bytes) { + return Ok(img); } - }; - Ok(Arc::new(clip)) - } + let cursor = std::io::Cursor::new(bytes); + let reader = image::ImageReader::new(cursor).with_guessed_format()?; + Ok(reader.decode()?) + }, + ) + .await + .map_err(MediaConnectorError::Blocking)? } enum DecodedVideoFrames { Images(Vec), Rgb(DecodedRgbVideo), + #[cfg(feature = "opencv-video")] + RgbStream(DecodedRgbFrameStream), } async fn decode_video_frames( bytes: Bytes, cfg: VideoFetchConfig, + runtime: Arc, +) -> Result { + #[cfg(not(feature = "opencv-video"))] + let _ = &runtime; + #[cfg(feature = "opencv-video")] + let input_bytes = bytes.len(); + match video_decode_backend_override() { + Some("ffmpeg") => decode_video_bytes_with_ffmpeg(bytes, cfg).await, + Some("opencv") => { + #[cfg(feature = "opencv-video")] + { + let opencv_bytes = bytes.clone(); + let opencv_runtime = runtime.clone(); + let result = task::spawn_blocking(move || { + decode_video_with_opencv_bytes_logged( + opencv_bytes, + input_bytes, + cfg, + &opencv_runtime, + ) + }) + .await + .map_err(MediaConnectorError::Blocking)?; + match result { + Ok(frames) => Ok(frames), + Err(error) => { + if log_video_decode_timing_enabled() { + info!( + error = %error, + "smg_mm_timing video_decode_opencv_buffer_fallback" + ); + } + decode_video_bytes_with_tempfile(bytes, cfg, runtime).await + } + } + } + #[cfg(not(feature = "opencv-video"))] + { + Err(MediaConnectorError::VideoDecode( + "SMG_VIDEO_DECODE_BACKEND=opencv requires the opencv-video feature".to_string(), + )) + } + } + Some(backend) => Err(MediaConnectorError::VideoDecode(format!( + "unsupported SMG_VIDEO_DECODE_BACKEND={backend}; expected auto, opencv, or ffmpeg" + ))), + None => { + #[cfg(feature = "opencv-video")] + { + let opencv_bytes = bytes.clone(); + let opencv_runtime = runtime.clone(); + let opencv_result = task::spawn_blocking(move || { + decode_video_with_opencv_bytes_logged( + opencv_bytes, + input_bytes, + cfg, + &opencv_runtime, + ) + }) + .await + .map_err(MediaConnectorError::Blocking)?; + match opencv_result { + Ok(frames) => Ok(frames), + Err(opencv_error) => { + if log_video_decode_timing_enabled() { + info!( + error = %opencv_error, + "smg_mm_timing video_decode_auto_opencv_fallback" + ); + } + decode_video_bytes_with_tempfile(bytes, cfg, runtime) + .await + .map_err(|fallback_error| { + MediaConnectorError::VideoDecode(format!( + "buffered OpenCV decode failed: {opencv_error}; tempfile fallback failed: {fallback_error}" + )) + }) + } + } + } + #[cfg(not(feature = "opencv-video"))] + { + decode_video_bytes_with_ffmpeg(bytes, cfg).await + } + } + } +} + +#[cfg(feature = "opencv-video")] +async fn decode_video_bytes_with_tempfile( + bytes: Bytes, + cfg: VideoFetchConfig, + runtime: Arc, +) -> Result { + let input_bytes = bytes.len(); + let input_file = { + let bytes = bytes.clone(); + task::spawn_blocking(move || write_temp_video_file(&bytes)) + .await + .map_err(MediaConnectorError::Blocking)?? + }; + decode_video_frames_from_path(input_file.path(), input_bytes, Some(&bytes), cfg, runtime).await +} + +async fn decode_video_bytes_with_ffmpeg( + bytes: Bytes, + cfg: VideoFetchConfig, ) -> Result { let input_bytes = bytes.len(); let input_file = { @@ -410,14 +552,27 @@ async fn decode_video_frames( .map_err(MediaConnectorError::Blocking)?? }; let input_path = input_file.path().to_path_buf(); + decode_video_with_ffmpeg(&input_path, input_bytes, Some(&bytes), cfg).await +} + +#[cfg(feature = "opencv-video")] +async fn decode_video_frames_from_path( + input_path: &std::path::Path, + input_bytes: usize, + input_data: Option<&Bytes>, + cfg: VideoFetchConfig, + runtime: Arc, +) -> Result { + #[cfg(not(feature = "opencv-video"))] + let _ = &runtime; match video_decode_backend_override() { - Some("ffmpeg") => decode_video_with_ffmpeg(&input_path, input_bytes, cfg).await, + Some("ffmpeg") => decode_video_with_ffmpeg(input_path, input_bytes, input_data, cfg).await, Some("opencv") => { #[cfg(feature = "opencv-video")] { - let input_path = input_path.clone(); + let input_path = input_path.to_path_buf(); task::spawn_blocking(move || { - decode_video_with_opencv_logged(&input_path, input_bytes, cfg) + decode_video_with_opencv_logged(&input_path, input_bytes, cfg, &runtime) }) .await .map_err(MediaConnectorError::Blocking)? @@ -437,9 +592,15 @@ async fn decode_video_frames( { // OpenCV samples by frame index while the FFmpeg fallback uses an // fps filter, so the fallback can select a different frame set. - let opencv_input_path = input_path.clone(); + let opencv_input_path = input_path.to_path_buf(); + let opencv_runtime = runtime.clone(); let opencv_result = task::spawn_blocking(move || { - decode_video_with_opencv_logged(&opencv_input_path, input_bytes, cfg) + decode_video_with_opencv_logged( + &opencv_input_path, + input_bytes, + cfg, + &opencv_runtime, + ) }) .await .map_err(MediaConnectorError::Blocking)?; @@ -454,7 +615,9 @@ async fn decode_video_frames( ); } - match decode_video_with_ffmpeg(&input_path, input_bytes, cfg).await { + match decode_video_with_ffmpeg(input_path, input_bytes, input_data, cfg) + .await + { Ok(frames) => Ok(frames), Err(ffmpeg_error) => Err(MediaConnectorError::VideoDecode(format!( "OpenCV decode failed: {opencv_error}; ffmpeg fallback failed: {ffmpeg_error}" @@ -466,22 +629,68 @@ async fn decode_video_frames( #[cfg(not(feature = "opencv-video"))] { - decode_video_with_ffmpeg(&input_path, input_bytes, cfg).await + decode_video_with_ffmpeg(input_path, input_bytes, input_data, cfg).await } } } } +fn validate_video_fetch_config(cfg: VideoFetchConfig) -> Result<(), MediaConnectorError> { + if cfg.max_frames == 0 { + return Err(MediaConnectorError::VideoDecode( + "max_frames must be greater than 0".to_string(), + )); + } + if cfg.min_frames == 0 { + return Err(MediaConnectorError::VideoDecode( + "min_frames must be greater than 0".to_string(), + )); + } + if cfg.min_frames > cfg.max_frames { + return Err(MediaConnectorError::VideoDecode( + "min_frames must be less than or equal to max_frames".to_string(), + )); + } + if !cfg.sample_fps.is_finite() || cfg.sample_fps <= 0.0 { + return Err(MediaConnectorError::VideoDecode( + "sample_fps must be finite and greater than 0".to_string(), + )); + } + Ok(()) +} + +fn video_clip_from_decoded( + decoded: DecodedVideoFrames, + bytes: Bytes, + source: VideoSource, + hash: String, +) -> VideoClip { + match decoded { + DecodedVideoFrames::Images(frames) => VideoClip::new(frames, bytes, source, hash), + DecodedVideoFrames::Rgb(rgb_video) => VideoClip::new_rgb(rgb_video, bytes, source, hash), + #[cfg(feature = "opencv-video")] + DecodedVideoFrames::RgbStream(stream) => { + VideoClip::new_rgb_stream(stream, bytes, source, hash) + } + } +} + #[cfg(feature = "opencv-video")] fn decode_video_with_opencv_logged( input_path: &std::path::Path, input_bytes: usize, cfg: VideoFetchConfig, + runtime: &MultimodalRuntime, ) -> Result { - let started = Instant::now(); - let result = decode_video_with_opencv_file(input_path, cfg); + let started = video_decode_timing_started(); + let result = decode_video_with_opencv_file_stream(input_path, cfg, runtime); + let backend = if matches!(&result, Ok(DecodedVideoFrames::RgbStream(_))) { + "opencv_stream_startup" + } else { + "opencv" + }; match &result { - Ok(_) => log_video_decode_backend_timing("opencv", started, input_bytes, cfg, None), + Ok(_) => log_video_decode_backend_timing(backend, started, input_bytes, cfg, None), Err(error) => { log_video_decode_backend_timing("opencv", started, input_bytes, cfg, Some(error)); } @@ -489,6 +698,35 @@ fn decode_video_with_opencv_logged( result } +#[cfg(feature = "opencv-video")] +fn decode_video_with_opencv_bytes_logged( + bytes: Bytes, + input_bytes: usize, + cfg: VideoFetchConfig, + runtime: &MultimodalRuntime, +) -> Result { + let started = video_decode_timing_started(); + let result = decode_video_with_opencv_bytes_stream(bytes, cfg, runtime); + let backend = if matches!(&result, Ok(DecodedVideoFrames::RgbStream(_))) { + "opencv_buffer_stream_startup" + } else { + "opencv_buffer" + }; + match &result { + Ok(_) => log_video_decode_backend_timing(backend, started, input_bytes, cfg, None), + Err(error) => { + log_video_decode_backend_timing( + "opencv_buffer", + started, + input_bytes, + cfg, + Some(error), + ); + } + } + result +} + fn video_decode_backend_override() -> Option<&'static str> { VIDEO_DECODE_BACKEND .get_or_init(|| { @@ -517,9 +755,13 @@ fn log_video_decode_timing_enabled() -> bool { }) } +fn video_decode_timing_started() -> Option { + log_video_decode_timing_enabled().then(Instant::now) +} + fn log_video_decode_backend_timing( backend: &str, - started: Instant, + started: Option, input_bytes: usize, cfg: VideoFetchConfig, error: Option<&MediaConnectorError>, @@ -527,7 +769,9 @@ fn log_video_decode_backend_timing( if !log_video_decode_timing_enabled() { return; } - let elapsed_ms = started.elapsed().as_secs_f64() * 1000.0; + let elapsed_ms = started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(); match error { Some(error) => info!( backend, @@ -554,9 +798,10 @@ fn log_video_decode_backend_timing( } #[cfg(feature = "opencv-video")] -fn decode_video_with_opencv_file( +fn decode_video_with_opencv_file_stream( input_path: &std::path::Path, cfg: VideoFetchConfig, + runtime: &MultimodalRuntime, ) -> Result { let input = input_path.to_str().ok_or_else(|| { MediaConnectorError::VideoDecode(format!( @@ -564,9 +809,215 @@ fn decode_video_with_opencv_file( input_path.display() )) })?; + let active_decode = runtime.enter_video_decode(OPENCV_DECODE_BURST_COALESCE); + let active_decodes = active_decode.count(); + let decoder_threads = + adaptive_opencv_decoder_threads(active_decode.available_parallelism(), active_decodes); + let capture = open_opencv_video_capture(input, decoder_threads)?; + // Request-level parallelism already overlaps decode and preprocessing. + // Adding another producer thread under load only increases CPU contention. + if active_decode.count() > 1 { + return decode_video_from_opencv_capture(capture, cfg); + } + decode_video_from_opencv_capture_stream(capture, cfg, active_decode) +} - let mut capture = open_opencv_video_capture(input)?; +#[cfg(feature = "opencv-video")] +fn decode_video_with_opencv_bytes_stream( + bytes: Bytes, + cfg: VideoFetchConfig, + runtime: &MultimodalRuntime, +) -> Result { + let active_decode = runtime.enter_video_decode(OPENCV_DECODE_BURST_COALESCE); + let active_decodes = active_decode.count(); + let decoder_threads = + adaptive_opencv_decoder_threads(active_decode.available_parallelism(), active_decodes); + let capture = open_opencv_video_capture_from_buffer(bytes, decoder_threads)?; + // Recheck after capture startup so concurrent requests have time to enter. + if active_decode.count() > 1 { + return decode_video_from_opencv_capture(capture, cfg); + } + decode_video_from_opencv_capture_stream(capture, cfg, active_decode) +} + +#[cfg(feature = "opencv-video")] +trait OpenCvCaptureOwner { + fn capture_mut(&mut self) -> &mut videoio::VideoCapture; +} +#[cfg(feature = "opencv-video")] +impl OpenCvCaptureOwner for videoio::VideoCapture { + fn capture_mut(&mut self) -> &mut videoio::VideoCapture { + self + } +} + +#[cfg(feature = "opencv-video")] +impl OpenCvCaptureOwner for crate::opencv_buffer::BufferedCapture { + fn capture_mut(&mut self) -> &mut videoio::VideoCapture { + self.capture_mut() + } +} + +#[cfg(feature = "opencv-video")] +fn decode_video_from_opencv_capture_stream( + mut capture: C, + cfg: VideoFetchConfig, + active_decode: ActiveVideoDecode, +) -> Result +where + C: OpenCvCaptureOwner + Send + 'static, +{ + let total_frames = capture + .capture_mut() + .get(videoio::CAP_PROP_FRAME_COUNT) + .map_err(opencv_decode_error)? + .round() + .max(0.0) as usize; + if total_frames == 0 { + return Err(MediaConnectorError::VideoDecode( + "OpenCV reported zero video frames".to_string(), + )); + } + let fps = capture + .capture_mut() + .get(videoio::CAP_PROP_FPS) + .map_err(opencv_decode_error)?; + let frame_indices = opencv_frame_indices(total_frames, fps, cfg); + if frame_indices.is_empty() { + return Err(MediaConnectorError::VideoDecode( + "OpenCV video sampling produced no frame indices".to_string(), + )); + } + + let expected_frames = frame_indices.len(); + let sampled_frame_counts = counted_frame_indices(&frame_indices); + let (sender, receiver) = sync_channel(2); + std::thread::Builder::new() + .name("smg-video-decode".to_string()) + .spawn(move || { + let _active_decode = active_decode; + if let Err(error) = decode_opencv_frames_to_stream( + capture.capture_mut(), + sampled_frame_counts, + expected_frames, + &sender, + ) { + let _ = sender.send(Err(error.to_string())); + } + }) + .map_err(MediaConnectorError::Io)?; + + Ok(DecodedVideoFrames::RgbStream(DecodedRgbFrameStream::new( + expected_frames, + receiver, + ))) +} + +#[cfg(feature = "opencv-video")] +fn decode_opencv_frames_to_stream( + capture: &mut videoio::VideoCapture, + sampled_frame_counts: Vec<(usize, usize)>, + expected_frames: usize, + sender: &SyncSender>, +) -> Result<(), MediaConnectorError> { + let timeout = video_process_timeout(); + let started = Instant::now(); + let mut decoded_pos: i64 = -1; + let mut emitted = 0usize; + let mut decoded_bytes = 0usize; + let mut bgr_frame = Mat::default(); + + for (idx, repeat_count) in sampled_frame_counts { + while decoded_pos + 1 < idx as i64 { + if started.elapsed() >= timeout { + return Err(MediaConnectorError::VideoDecode(format!( + "OpenCV timed out after {:.3} seconds", + timeout.as_secs_f64() + ))); + } + if !capture.grab().map_err(opencv_decode_error)? { + return Err(MediaConnectorError::VideoDecode(format!( + "OpenCV could not grab intervening frame to reach sampled frame {idx}" + ))); + } + decoded_pos += 1; + } + + if started.elapsed() >= timeout { + return Err(MediaConnectorError::VideoDecode(format!( + "OpenCV timed out after {:.3} seconds", + timeout.as_secs_f64() + ))); + } + let read_successful = capture.read(&mut bgr_frame).map_err(opencv_decode_error)?; + if started.elapsed() >= timeout { + return Err(MediaConnectorError::VideoDecode(format!( + "OpenCV timed out after {:.3} seconds", + timeout.as_secs_f64() + ))); + } + decoded_pos = idx as i64; + if !read_successful || bgr_frame.empty() { + continue; + } + let width = u32::try_from(bgr_frame.cols()).map_err(|_| { + MediaConnectorError::VideoDecode(format!( + "OpenCV produced invalid BGR frame width: {}", + bgr_frame.cols() + )) + })?; + let height = u32::try_from(bgr_frame.rows()).map_err(|_| { + MediaConnectorError::VideoDecode(format!( + "OpenCV produced invalid BGR frame height: {}", + bgr_frame.rows() + )) + })?; + let frame_size = rawvideo_frame_size(width, height)?; + decoded_bytes = decoded_bytes.checked_add(frame_size).ok_or_else(|| { + MediaConnectorError::VideoDecode( + "decoded video byte size overflow while streaming frames".to_string(), + ) + })?; + ensure_decoded_byte_limit(decoded_bytes)?; + let bgr_bytes = bgr_frame.data_bytes().map_err(opencv_decode_error)?; + if bgr_bytes.len() < frame_size { + return Err(MediaConnectorError::VideoDecode(format!( + "OpenCV produced {} BGR bytes for {width}x{height} frame, expected {frame_size}", + bgr_bytes.len() + ))); + } + let frame = OwnedRgbFrame { + width, + height, + data: Bytes::copy_from_slice(&bgr_bytes[..frame_size]), + channel_order: RgbChannelOrder::Bgr, + }; + for _ in 0..repeat_count { + if sender.send(Ok(frame.clone())).is_err() { + return Ok(()); + } + emitted += 1; + } + } + + if emitted != expected_frames { + return Err(MediaConnectorError::VideoDecode(format!( + "OpenCV produced {emitted} sampled frames, expected {expected_frames}" + ))); + } + Ok(()) +} + +#[cfg(feature = "opencv-video")] +fn decode_video_from_opencv_capture( + mut capture: C, + cfg: VideoFetchConfig, +) -> Result +where + C: OpenCvCaptureOwner, +{ + let capture = capture.capture_mut(); let total_frames = capture .get(videoio::CAP_PROP_FRAME_COUNT) .map_err(opencv_decode_error)? @@ -589,6 +1040,7 @@ fn decode_video_with_opencv_file( } let sampled_frame_counts = counted_frame_indices(&frame_indices); + let unique_sampled_frames = sampled_frame_counts.len(); let mut data = Vec::new(); let mut frames = Vec::new(); frames.try_reserve(frame_indices.len()).map_err(|e| { @@ -602,8 +1054,15 @@ fn decode_video_with_opencv_file( let timeout = video_process_timeout(); let started = Instant::now(); - // Seek directly to sampled frames instead of scanning every intervening - // frame, which can be prohibitively slow for long clips. + // Advance to each sampled frame by SEQUENTIALLY grabbing the intervening frames + // (cheap decode-without-retrieve) and `read`ing only the sampled ones, instead of + // calling `set(CAP_PROP_POS_FRAMES)` per frame. OpenCV's POS_FRAMES set flushes/ + // re-seeks the decoder on every call (~10 ms/frame even for adjacent frames); + // sequential grab is ~1-2 ms/frame. This is verified against the old + // per-frame seek on both dense and sparse (non-keyframe) sampling, so + // accuracy is unchanged. `sampled_frame_counts` is monotonic. + // Index of the most recently decoded frame (-1 = nothing read yet). + let mut decoded_pos: i64 = -1; for (idx, repeat_count) in sampled_frame_counts { if started.elapsed() >= timeout { return Err(MediaConnectorError::VideoDecode(format!( @@ -612,16 +1071,26 @@ fn decode_video_with_opencv_file( ))); } - if !capture - .set(videoio::CAP_PROP_POS_FRAMES, idx as f64) - .map_err(opencv_decode_error)? - { - return Err(MediaConnectorError::VideoDecode(format!( - "OpenCV could not seek to sampled frame {idx}" - ))); + // Skip-decode the frames between the current position and `idx` so the + // following `read` lands on `idx` without a decoder flush/seek. + while decoded_pos + 1 < idx as i64 { + if started.elapsed() >= timeout { + return Err(MediaConnectorError::VideoDecode(format!( + "OpenCV timed out after {:.3} seconds", + timeout.as_secs_f64() + ))); + } + if !capture.grab().map_err(opencv_decode_error)? { + return Err(MediaConnectorError::VideoDecode(format!( + "OpenCV could not grab intervening frame to reach sampled frame {idx}" + ))); + } + decoded_pos += 1; } - if !capture.read(&mut bgr_frame).map_err(opencv_decode_error)? || bgr_frame.empty() { + let read_successful = capture.read(&mut bgr_frame).map_err(opencv_decode_error)?; + decoded_pos = idx as i64; + if !read_successful || bgr_frame.empty() { continue; } @@ -641,6 +1110,14 @@ fn decode_video_with_opencv_file( )) })?; let frame_size = rawvideo_frame_size(decoded_width, decoded_height)?; + if data.capacity() == 0 { + let decoded_bytes = checked_decoded_rgb_bytes(unique_sampled_frames, frame_size)?; + data.try_reserve_exact(decoded_bytes).map_err(|e| { + MediaConnectorError::VideoDecode(format!( + "failed to reserve {decoded_bytes} decoded video bytes: {e}" + )) + })?; + } let rgb_bytes = rgb_frame.data_bytes().map_err(opencv_decode_error)?; if rgb_bytes.len() < frame_size { return Err(MediaConnectorError::VideoDecode(format!( @@ -648,20 +1125,15 @@ fn decode_video_with_opencv_file( rgb_bytes.len() ))); } + let new_len = data.len().checked_add(frame_size).ok_or_else(|| { + MediaConnectorError::VideoDecode(format!( + "decoded video byte size overflow while appending {frame_size} bytes" + )) + })?; + ensure_decoded_byte_limit(new_len)?; + let offset = data.len(); + data.extend_from_slice(&rgb_bytes[..frame_size]); for _ in 0..repeat_count { - let new_len = data.len().checked_add(frame_size).ok_or_else(|| { - MediaConnectorError::VideoDecode(format!( - "decoded video byte size overflow while appending {frame_size} bytes" - )) - })?; - ensure_decoded_byte_limit(new_len)?; - data.try_reserve(frame_size).map_err(|e| { - MediaConnectorError::VideoDecode(format!( - "failed to reserve {frame_size} decoded video bytes: {e}" - )) - })?; - let offset = data.len(); - data.extend_from_slice(&rgb_bytes[..frame_size]); frames.push(DecodedRgbFrame { width: decoded_width, height: decoded_height, @@ -691,17 +1163,36 @@ fn decode_video_with_opencv_file( } #[cfg(feature = "opencv-video")] -fn open_opencv_video_capture(input: &str) -> Result { - let capture = videoio::VideoCapture::from_file(input, videoio::CAP_FFMPEG) - .map_err(opencv_decode_error)?; - if capture.is_opened().map_err(opencv_decode_error)? { - return Ok(capture); +fn open_opencv_video_capture_from_buffer( + bytes: Bytes, + decoder_threads: i32, +) -> Result { + crate::opencv_buffer::open_capture(bytes, decoder_threads).map_err(|error| { + MediaConnectorError::VideoDecode(format!("OpenCV could not open video buffer: {error}")) + }) +} + +#[cfg(feature = "opencv-video")] +fn open_opencv_video_capture( + input: &str, + decoder_threads: i32, +) -> Result { + let params = Vector::from_slice(&[videoio::CAP_PROP_N_THREADS, decoder_threads]); + if let Ok(capture) = + videoio::VideoCapture::from_file_with_params(input, videoio::CAP_FFMPEG, ¶ms) + { + if capture.is_opened().map_err(opencv_decode_error)? { + return Ok(capture); + } } - let capture = - videoio::VideoCapture::from_file(input, videoio::CAP_ANY).map_err(opencv_decode_error)?; - if capture.is_opened().map_err(opencv_decode_error)? { - return Ok(capture); + for backend in [videoio::CAP_FFMPEG, videoio::CAP_ANY] { + let Ok(capture) = videoio::VideoCapture::from_file(input, backend) else { + continue; + }; + if capture.is_opened().map_err(opencv_decode_error)? { + return Ok(capture); + } } Err(MediaConnectorError::VideoDecode(format!( @@ -709,6 +1200,38 @@ fn open_opencv_video_capture(input: &str) -> Result i32 { + let available_cpus = available_cpus.max(1); + let active_decodes = active_decodes.max(1); + + // Once eight or more independent decoders fill the CPU quota, codec-level + // threading only adds scheduler contention. + if active_decodes >= OPENCV_LOW_CONCURRENCY_LIMIT && active_decodes >= available_cpus { + return 1; + } + + let (decoder_budget, max_threads) = if active_decodes <= OPENCV_LOW_CONCURRENCY_LIMIT { + let max_threads = if active_decodes <= 2 { 16 } else { 8 }; + ( + available_cpus.saturating_mul(OPENCV_LOW_CONCURRENCY_CPU_MULTIPLIER), + max_threads, + ) + } else { + // Independent decoders supply request-level parallelism at high + // concurrency. Reserve roughly one seventh of the CPU quota for frame + // copies, request handling, and other non-decoder work. + ( + available_cpus + .saturating_mul(OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_NUMERATOR) + .div_ceil(OPENCV_HIGH_CONCURRENCY_CPU_BUDGET_DENOMINATOR), + MAX_OPENCV_DECODER_THREADS, + ) + }; + + (decoder_budget.max(1) / active_decodes).clamp(1, max_threads) as i32 +} + #[cfg(feature = "opencv-video")] fn opencv_frame_indices(total_frames: usize, fps: f64, cfg: VideoFetchConfig) -> Vec { let mut target_frames = if fps.is_finite() && fps > 0.0 { @@ -753,27 +1276,30 @@ fn opencv_decode_error(err: opencv::Error) -> MediaConnectorError { async fn decode_video_with_ffmpeg( input_path: &std::path::Path, input_bytes: usize, + input_data: Option<&Bytes>, cfg: VideoFetchConfig, ) -> Result { - if let Ok(metadata) = probe_video_metadata(input_path).await { - let started = Instant::now(); - match decode_video_with_ffmpeg_ppm(input_path, cfg, metadata).await { - Ok(rgb_video) => { - log_video_decode_backend_timing("ffmpeg_ppm_file", started, input_bytes, cfg, None); - return Ok(DecodedVideoFrames::Rgb(rgb_video)); - } - Err(error) => { - log_video_decode_backend_timing( - "ffmpeg_ppm_file", - started, - input_bytes, - cfg, - Some(&error), - ); - } + let duration_seconds = video_duration_seconds_for_input(input_path, input_data).await; + + let started = video_decode_timing_started(); + match decode_video_with_ffmpeg_ppm(input_path, cfg, duration_seconds).await { + Ok(rgb_video) => { + log_video_decode_backend_timing("ffmpeg_ppm_file", started, input_bytes, cfg, None); + return Ok(DecodedVideoFrames::Rgb(rgb_video)); + } + Err(error) => { + log_video_decode_backend_timing( + "ffmpeg_ppm_file", + started, + input_bytes, + cfg, + Some(&error), + ); } + } - let started = Instant::now(); + if let Ok(metadata) = probe_video_metadata(input_path).await { + let started = video_decode_timing_started(); match decode_video_with_ffmpeg_raw(input_path, cfg, metadata).await { Ok(rgb_video) => { log_video_decode_backend_timing("ffmpeg_raw_file", started, input_bytes, cfg, None); @@ -791,8 +1317,8 @@ async fn decode_video_with_ffmpeg( } } - let started = Instant::now(); - match decode_video_with_ffmpeg_png(input_path, cfg).await { + let started = video_decode_timing_started(); + match decode_video_with_ffmpeg_png(input_path, cfg, duration_seconds).await { Ok(frames) => { log_video_decode_backend_timing("ffmpeg_png_file", started, input_bytes, cfg, None); Ok(DecodedVideoFrames::Images(frames)) @@ -811,7 +1337,7 @@ async fn decode_video_with_ffmpeg( } fn write_temp_video_file(bytes: &[u8]) -> Result { - let started = Instant::now(); + let started = video_decode_timing_started(); let mut input_file = tempfile::Builder::new() .prefix("smg-video-") .suffix(video_temp_suffix(bytes)) @@ -821,7 +1347,9 @@ fn write_temp_video_file(bytes: &[u8]) -> Result Result { + let child = spawn_video_command(command, program)?; + let timeout = video_process_timeout(); + match time::timeout(timeout, child.wait_with_output()).await { + Ok(Ok(output)) => Ok(output), + Ok(Err(error)) => Err(MediaConnectorError::Io(error)), + Err(_) => Err(MediaConnectorError::VideoDecode(format!( + "{program} timed out after {:.3} seconds", + timeout.as_secs_f64() + ))), + } +} + +async fn run_video_command_output_with_stdout_capacity( + command: Command, program: &'static str, + stdout_capacity: usize, ) -> Result { + let child = spawn_video_command(command, program)?; + let timeout = video_process_timeout(); + match time::timeout( + timeout, + collect_video_command_output(child, stdout_capacity), + ) + .await + { + Ok(Ok(output)) => Ok(output), + Ok(Err(error)) => Err(MediaConnectorError::Io(error)), + Err(_) => Err(MediaConnectorError::VideoDecode(format!( + "{program} timed out after {:.3} seconds", + timeout.as_secs_f64() + ))), + } +} + +fn spawn_video_command( + mut command: Command, + program: &'static str, +) -> Result { command .stdout(Stdio::piped()) .stderr(Stdio::piped()) .kill_on_drop(true); - let child = command.spawn().map_err(|e| { + command.spawn().map_err(|e| { if e.kind() == std::io::ErrorKind::NotFound { MediaConnectorError::VideoDecode(format!( - "{program} executable not found; install {program} to decode video_url inputs" + "{program} executable not found; install {program} to decode video inputs" )) } else { MediaConnectorError::Io(e) } - })?; + }) +} - let timeout = video_process_timeout(); - match time::timeout(timeout, child.wait_with_output()).await { - Ok(Ok(output)) => Ok(output), - Ok(Err(error)) => Err(MediaConnectorError::Io(error)), - Err(_) => Err(MediaConnectorError::VideoDecode(format!( - "{program} timed out after {:.3} seconds", - timeout.as_secs_f64() - ))), - } +async fn collect_video_command_output( + mut child: tokio::process::Child, + stdout_capacity: usize, +) -> std::io::Result { + let mut stdout_pipe = child + .stdout + .take() + .ok_or_else(|| std::io::Error::other("video command stdout was not piped"))?; + let mut stderr_pipe = child + .stderr + .take() + .ok_or_else(|| std::io::Error::other("video command stderr was not piped"))?; + + let stdout = async move { + let mut bytes = Vec::with_capacity(stdout_capacity); + stdout_pipe.read_to_end(&mut bytes).await?; + Ok::<_, std::io::Error>(bytes) + }; + let stderr = async move { + let mut bytes = Vec::new(); + stderr_pipe.read_to_end(&mut bytes).await?; + Ok::<_, std::io::Error>(bytes) + }; + let status = child.wait(); + + let (stdout, stderr, status) = tokio::try_join!(stdout, stderr, status)?; + Ok(Output { + status, + stdout, + stderr, + }) } async fn decode_video_with_ffmpeg_ppm( input_path: &std::path::Path, cfg: VideoFetchConfig, - metadata: VideoMetadata, + duration_seconds: Option, ) -> Result { - let fps_filter = fps_filter_for_metadata(metadata, cfg); + let fps_filter = fps_filter_for_optional_duration(duration_seconds, cfg); let max_frames = cfg.max_frames.to_string(); - let frame_size = rawvideo_frame_size(metadata.width, metadata.height)?; - let target_frames = expected_sampled_frame_count(metadata, cfg); - let decoded_bytes = checked_decoded_rgb_bytes(target_frames, frame_size)?; - let output_limit = decoded_bytes - .checked_add(target_frames.saturating_mul(64)) - .unwrap_or_else(video_max_decoded_bytes) - .min(video_max_decoded_bytes()) - .to_string(); + let output_limit = video_max_decoded_bytes().to_string(); let mut command = Command::new("ffmpeg"); command - .args(["-hide_banner", "-loglevel", "error", "-nostdin", "-i"]) + .args([ + "-hide_banner", + "-loglevel", + "error", + "-nostdin", + "-threads", + "1", + "-i", + ]) .arg(input_path) .args([ + "-map", + "0:v:0", + "-an", + "-sn", + "-dn", "-vf", &fps_filter, "-frames:v", @@ -955,7 +1550,7 @@ async fn decode_video_with_ffmpeg_ppm( "rgb24", "pipe:1", ]); - let output = run_video_command_output(command, "ffmpeg").await?; + let output = run_video_command_output_with_stdout_capacity(command, "ffmpeg", 0).await?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); @@ -988,11 +1583,18 @@ async fn decode_video_with_ffmpeg_raw( "-loglevel", "error", "-nostdin", + "-threads", + "1", "-noautorotate", "-i", ]) .arg(input_path) .args([ + "-map", + "0:v:0", + "-an", + "-sn", + "-dn", "-vf", &fps_filter, "-frames:v", @@ -1005,7 +1607,12 @@ async fn decode_video_with_ffmpeg_raw( "rgb24", "pipe:1", ]); - let output = run_video_command_output(command, "ffmpeg").await?; + let output = run_video_command_output_with_stdout_capacity( + command, + "ffmpeg", + video_stdout_prealloc_capacity(metadata, decoded_bytes), + ) + .await?; if !output.status.success() { let stderr = String::from_utf8_lossy(&output.stderr); @@ -1047,15 +1654,29 @@ async fn decode_video_with_ffmpeg_raw( async fn decode_video_with_ffmpeg_png( input_path: &std::path::Path, cfg: VideoFetchConfig, + duration_seconds: Option, ) -> Result, MediaConnectorError> { - let fps_filter = fps_filter_for_video(input_path, cfg).await; + let fps_filter = fps_filter_for_optional_duration(duration_seconds, cfg); let max_frames = cfg.max_frames.to_string(); let output_limit = video_max_decoded_bytes().to_string(); let mut command = Command::new("ffmpeg"); command - .args(["-hide_banner", "-loglevel", "error", "-nostdin", "-i"]) + .args([ + "-hide_banner", + "-loglevel", + "error", + "-nostdin", + "-threads", + "1", + "-i", + ]) .arg(input_path) .args([ + "-map", + "0:v:0", + "-an", + "-sn", + "-dn", "-vf", &fps_filter, "-frames:v", @@ -1112,7 +1733,6 @@ async fn probe_video_metadata( .args([ "-v", "error", - "-nostdin", "-select_streams", "v:0", "-show_entries", @@ -1160,7 +1780,14 @@ async fn probe_video_metadata( } fn fps_filter_for_metadata(metadata: VideoMetadata, cfg: VideoFetchConfig) -> String { - if let Some(duration) = metadata.duration_seconds { + fps_filter_for_optional_duration(metadata.duration_seconds, cfg) +} + +fn fps_filter_for_optional_duration( + duration_seconds: Option, + cfg: VideoFetchConfig, +) -> String { + if let Some(duration) = duration_seconds { if let Some(filter) = fps_filter_for_duration(duration, cfg) { return filter; } @@ -1180,6 +1807,13 @@ fn expected_sampled_frame_count(metadata: VideoMetadata, cfg: VideoFetchConfig) cfg.max_frames } +fn video_stdout_prealloc_capacity(metadata: VideoMetadata, expected_bytes: usize) -> usize { + match metadata.duration_seconds { + Some(duration) if duration.is_finite() && duration > 0.0 => expected_bytes, + _ => 0, + } +} + fn fps_filter_for_duration(duration: f64, cfg: VideoFetchConfig) -> Option { if !duration.is_finite() || duration <= 0.0 { return None; @@ -1191,14 +1825,17 @@ fn fps_filter_for_duration(duration: f64, cfg: VideoFetchConfig) -> Option String { - if let Ok(duration) = probe_video_duration_seconds(input_path).await { - if let Some(filter) = fps_filter_for_duration(duration, cfg) { - return filter; +async fn video_duration_seconds_for_input( + input_path: &std::path::Path, + input_data: Option<&Bytes>, +) -> Option { + if let Some(bytes) = input_data { + if let Some(duration) = parse_mp4_duration_seconds(bytes.as_ref()) { + return Some(duration); } } - format!("fps={}", cfg.sample_fps) + probe_video_duration_seconds(input_path).await.ok() } async fn probe_video_duration_seconds( @@ -1209,7 +1846,6 @@ async fn probe_video_duration_seconds( .args([ "-v", "error", - "-nostdin", "-show_entries", "format=duration", "-of", @@ -1253,6 +1889,88 @@ fn parse_ffmpeg_duration_seconds(stderr: &str) -> Option { Some(hours * 3600.0 + minutes * 60.0 + seconds) } +fn parse_mp4_duration_seconds(bytes: &[u8]) -> Option { + let mut pos = 0usize; + while let Some((kind, payload_start, payload_end)) = read_mp4_box(bytes, pos, bytes.len()) { + if kind == *b"moov" { + return parse_mp4_moov_duration_seconds(bytes, payload_start, payload_end); + } + pos = payload_end; + } + None +} + +fn parse_mp4_moov_duration_seconds(bytes: &[u8], start: usize, end: usize) -> Option { + let mut pos = start; + while let Some((kind, payload_start, payload_end)) = read_mp4_box(bytes, pos, end) { + if kind == *b"mvhd" { + return parse_mp4_mvhd_duration_seconds(&bytes[payload_start..payload_end]); + } + pos = payload_end; + } + None +} + +fn read_mp4_box(bytes: &[u8], pos: usize, limit: usize) -> Option<([u8; 4], usize, usize)> { + if pos.checked_add(8)? > limit || limit > bytes.len() { + return None; + } + let size32 = u32::from_be_bytes(bytes[pos..pos + 4].try_into().ok()?); + let kind: [u8; 4] = bytes[pos + 4..pos + 8].try_into().ok()?; + let mut header_len = 8usize; + let size = match size32 { + 0 => limit.checked_sub(pos)?, + 1 => { + if pos.checked_add(16)? > limit { + return None; + } + header_len = 16; + usize::try_from(u64::from_be_bytes( + bytes[pos + 8..pos + 16].try_into().ok()?, + )) + .ok()? + } + size => size as usize, + }; + if size < header_len { + return None; + } + let payload_start = pos.checked_add(header_len)?; + let payload_end = pos.checked_add(size)?; + if payload_end > limit || payload_start > payload_end { + return None; + } + Some((kind, payload_start, payload_end)) +} + +fn parse_mp4_mvhd_duration_seconds(payload: &[u8]) -> Option { + let version = *payload.first()?; + let (timescale_offset, duration_offset, duration_len) = match version { + 0 => (12usize, 16usize, 4usize), + 1 => (20usize, 24usize, 8usize), + _ => return None, + }; + let timescale_end = timescale_offset.checked_add(4)?; + let duration_end = duration_offset.checked_add(duration_len)?; + if duration_end > payload.len() || timescale_end > payload.len() { + return None; + } + let timescale = u32::from_be_bytes(payload[timescale_offset..timescale_end].try_into().ok()?); + if timescale == 0 { + return None; + } + let duration = if duration_len == 4 { + u32::from_be_bytes(payload[duration_offset..duration_end].try_into().ok()?) as f64 + } else { + u64::from_be_bytes(payload[duration_offset..duration_end].try_into().ok()?) as f64 + }; + let seconds = duration / timescale as f64; + seconds + .is_finite() + .then_some(seconds) + .filter(|value| *value > 0.0) +} + fn split_png_stream(bytes: &[u8]) -> Result, MediaConnectorError> { const PNG_SIG: &[u8; 8] = b"\x89PNG\r\n\x1a\n"; const IEND: &[u8; 4] = b"IEND"; @@ -1475,13 +2193,14 @@ fn skip_ppm_whitespace_and_comments(bytes: &[u8], pos: &mut usize) { #[cfg(test)] mod tests { use super::{ - parse_ffmpeg_duration_seconds, parse_ppm_stream, split_png_stream, video_temp_suffix, + parse_ffmpeg_duration_seconds, parse_mp4_duration_seconds, parse_ppm_stream, + split_png_stream, video_stdout_prealloc_capacity, video_temp_suffix, VideoMetadata, }; const TINY_PNG: &[u8] = &[ 137, 80, 78, 71, 13, 10, 26, 10, 0, 0, 0, 13, 73, 72, 68, 82, 0, 0, 0, 1, 0, 0, 0, 1, 8, 4, 0, 0, 0, 181, 28, 12, 2, 0, 0, 0, 11, 73, 68, 65, 84, 120, 218, 99, 96, 96, 0, 0, 0, 3, 0, - 1, 43, 9, 141, 84, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, + 1, 43, 9, 77, 132, 0, 0, 0, 0, 73, 69, 78, 68, 174, 66, 96, 130, ]; #[test] @@ -1505,6 +2224,54 @@ mod tests { assert_eq!(parse_ffmpeg_duration_seconds(stderr), Some(83.45)); } + fn mp4_box(kind: [u8; 4], payload: Vec) -> Vec { + let size = u32::try_from(payload.len() + 8).unwrap(); + let mut bytes = Vec::with_capacity(size as usize); + bytes.extend_from_slice(&size.to_be_bytes()); + bytes.extend_from_slice(&kind); + bytes.extend_from_slice(&payload); + bytes + } + + #[test] + fn parses_mp4_mvhd_v0_duration() { + let mut mvhd = vec![0; 20]; + mvhd[12..16].copy_from_slice(&1_000u32.to_be_bytes()); + mvhd[16..20].copy_from_slice(&2_000u32.to_be_bytes()); + let mut file = mp4_box(*b"ftyp", b"isom".to_vec()); + file.extend_from_slice(&mp4_box(*b"moov", mp4_box(*b"mvhd", mvhd))); + + assert_eq!(parse_mp4_duration_seconds(&file), Some(2.0)); + } + + #[test] + fn parses_mp4_mvhd_v1_duration() { + let mut mvhd = vec![0; 32]; + mvhd[0] = 1; + mvhd[20..24].copy_from_slice(&90_000u32.to_be_bytes()); + mvhd[24..32].copy_from_slice(&180_000u64.to_be_bytes()); + let file = mp4_box(*b"moov", mp4_box(*b"mvhd", mvhd)); + + assert_eq!(parse_mp4_duration_seconds(&file), Some(2.0)); + } + + #[test] + fn preallocates_video_stdout_only_with_known_duration() { + let known = VideoMetadata { + width: 16, + height: 16, + duration_seconds: Some(1.0), + }; + let unknown = VideoMetadata { + width: 16, + height: 16, + duration_seconds: None, + }; + + assert_eq!(video_stdout_prealloc_capacity(known, 4096), 4096); + assert_eq!(video_stdout_prealloc_capacity(unknown, 4096), 0); + } + #[test] fn detects_video_temp_suffix_from_container_header() { let mut mp4 = vec![0; 12]; @@ -1565,4 +2332,21 @@ mod tests { let indices = super::opencv_frame_indices(1, 30.0, cfg); assert_eq!(indices, vec![0, 0, 0, 0]); } + + #[cfg(feature = "opencv-video")] + #[test] + fn opencv_decoder_threads_share_cpu_budget_across_active_decodes() { + assert_eq!(super::adaptive_opencv_decoder_threads(224, 1), 16); + assert_eq!(super::adaptive_opencv_decoder_threads(2, 1), 4); + assert_eq!(super::adaptive_opencv_decoder_threads(4, 2), 4); + assert_eq!(super::adaptive_opencv_decoder_threads(8, 4), 4); + assert_eq!(super::adaptive_opencv_decoder_threads(8, 8), 1); + assert_eq!(super::adaptive_opencv_decoder_threads(8, 9), 1); + assert_eq!(super::adaptive_opencv_decoder_threads(16, 8), 4); + assert_eq!(super::adaptive_opencv_decoder_threads(16, 16), 1); + assert_eq!(super::adaptive_opencv_decoder_threads(224, 8), 8); + assert_eq!(super::adaptive_opencv_decoder_threads(224, 32), 6); + assert_eq!(super::adaptive_opencv_decoder_threads(8, 32), 1); + assert_eq!(super::adaptive_opencv_decoder_threads(1, 0), 2); + } } diff --git a/crates/multimodal/src/opencv_buffer.rs b/crates/multimodal/src/opencv_buffer.rs new file mode 100644 index 000000000..7db38b6af --- /dev/null +++ b/crates/multimodal/src/opencv_buffer.rs @@ -0,0 +1,55 @@ +//! Safe wrapper for OpenCV's buffered video capture constructor. +#![allow(unsafe_code)] + +use std::ffi::{c_char, c_void, CStr}; + +use bytes::Bytes; +use opencv::{traits::OpenCVFromExtern, videoio}; + +unsafe extern "C" { + fn smg_opencv_capture_from_buffer( + data: *const u8, + size: usize, + decoder_threads: i32, + error: *mut c_char, + error_capacity: usize, + ) -> *mut c_void; +} + +pub(crate) struct BufferedCapture { + capture: videoio::VideoCapture, + _bytes: Bytes, +} + +impl BufferedCapture { + pub(crate) fn capture_mut(&mut self) -> &mut videoio::VideoCapture { + &mut self.capture + } +} + +pub(crate) fn open_capture(bytes: Bytes, decoder_threads: i32) -> Result { + let mut error = [0 as c_char; 512]; + // SAFETY: `BufferedCapture` owns `bytes` for at least as long as the capture. + let capture = unsafe { + smg_opencv_capture_from_buffer( + bytes.as_ptr(), + bytes.len(), + decoder_threads, + error.as_mut_ptr(), + error.len(), + ) + }; + if capture.is_null() { + // SAFETY: the bridge always writes a NUL-terminated message on failure. + return Err(unsafe { CStr::from_ptr(error.as_ptr()) } + .to_string_lossy() + .into_owned()); + } + + Ok(BufferedCapture { + // SAFETY: the bridge returns a heap-allocated cv::VideoCapture compatible + // with the opencv crate's generated ownership wrapper. + capture: unsafe { videoio::VideoCapture::opencv_from_extern(capture) }, + _bytes: bytes, + }) +} diff --git a/crates/multimodal/src/opencv_buffer_capture.cpp b/crates/multimodal/src/opencv_buffer_capture.cpp new file mode 100644 index 000000000..c925a6e2f --- /dev/null +++ b/crates/multimodal/src/opencv_buffer_capture.cpp @@ -0,0 +1,96 @@ +#include +#include + +#include +#include +#include +#include +#include +#include + +namespace { + +class MemoryStreamReader final : public cv::IStreamReader { + public: + MemoryStreamReader(const uint8_t* data, size_t size) + : data_(data), size_(size) {} + + long long read(char* buffer, long long size) override { + if (size <= 0 || position_ >= size_) { + return 0; + } + const size_t count = + std::min(static_cast(size), size_ - position_); + std::memcpy(buffer, data_ + position_, count); + position_ += count; + return static_cast(count); + } + + long long seek(long long offset, int origin) override { + long long base = 0; + if (origin == SEEK_CUR) { + base = static_cast(position_); + } else if (origin == SEEK_END) { + base = static_cast(size_); + } else if (origin != SEEK_SET) { + return -1; + } + + const long long next = base + offset; + if (next < 0 || static_cast(next) > size_) { + return -1; + } + position_ = static_cast(next); + return next; + } + + private: + const uint8_t* data_; + size_t size_; + size_t position_ = 0; +}; + +void set_error(char* output, size_t capacity, const char* message) { + if (output == nullptr || capacity == 0) { + return; + } + std::snprintf(output, capacity, "%s", message); +} + +} // namespace + +extern "C" void* smg_opencv_capture_from_buffer(const uint8_t* data, + size_t size, + int decoder_threads, + char* error, + size_t error_capacity) { + try { + if (data == nullptr || size == 0) { + set_error(error, error_capacity, "video buffer is empty"); + return nullptr; + } + + for (const auto backend : + cv::videoio_registry::getStreamBufferedBackends()) { + if (!cv::videoio_registry::hasBackend(backend)) { + continue; + } + cv::Ptr reader = + cv::makePtr(data, size); + auto* capture = new cv::VideoCapture( + reader, static_cast(backend), + std::vector{cv::CAP_PROP_N_THREADS, decoder_threads}); + if (capture->isOpened()) { + return capture; + } + delete capture; + } + set_error(error, error_capacity, + "OpenCV has no usable buffered video backend"); + } catch (const std::exception& exception) { + set_error(error, error_capacity, exception.what()); + } catch (...) { + set_error(error, error_capacity, "unknown OpenCV buffered capture error"); + } + return nullptr; +} diff --git a/crates/multimodal/src/registry/kimi_k25.rs b/crates/multimodal/src/registry/kimi_k25.rs index 423bf0cf0..05a5ad981 100644 --- a/crates/multimodal/src/registry/kimi_k25.rs +++ b/crates/multimodal/src/registry/kimi_k25.rs @@ -89,7 +89,7 @@ impl ModelProcessorSpec for KimiK25VisionSpec { ]) } - fn keep_on_cpu_keys(&self) -> Vec { + fn cpu_resident_tensor_keys(&self) -> Vec { vec!["grid_thws".to_string()] } } diff --git a/crates/multimodal/src/registry/mod.rs b/crates/multimodal/src/registry/mod.rs index e9a0e9108..850080a24 100644 --- a/crates/multimodal/src/registry/mod.rs +++ b/crates/multimodal/src/registry/mod.rs @@ -79,7 +79,7 @@ pub(super) mod test_helpers { use crate::{ types::ImageSize, - vision::processor::{ModelSpecificValue, PreprocessedEncoderInputs}, + vision::processor::{EncoderInput, ModelSpecificValue, PreprocessedEncoderInputs}, }; pub struct TestTokenizer { @@ -160,7 +160,7 @@ pub(super) mod test_helpers { ) -> PreprocessedEncoderInputs { let sizes: Vec<(u32, u32)> = item_sizes.iter().map(|s| (s.height, s.width)).collect(); PreprocessedEncoderInputs { - encoder_input: ndarray::ArrayD::zeros(vec![1, 3, 336, 336]), + encoder_input: EncoderInput::Dense(ndarray::ArrayD::zeros(vec![1, 3, 336, 336])), feature_token_counts: feature_token_counts.to_vec(), item_sizes: sizes, model_specific: HashMap::new(), @@ -187,7 +187,7 @@ pub(super) mod test_helpers { }, ); PreprocessedEncoderInputs { - encoder_input: ndarray::ArrayD::zeros(vec![1, 3, 336, 336]), + encoder_input: EncoderInput::Dense(ndarray::ArrayD::zeros(vec![1, 3, 336, 336])), feature_token_counts: vec![0; sizes.len()], item_sizes: sizes, model_specific, diff --git a/crates/multimodal/src/registry/qwen3_vl.rs b/crates/multimodal/src/registry/qwen3_vl.rs index 4cbd49c9b..a124d1449 100644 --- a/crates/multimodal/src/registry/qwen3_vl.rs +++ b/crates/multimodal/src/registry/qwen3_vl.rs @@ -287,7 +287,7 @@ impl ModelProcessorSpec for Qwen3VLVisionSpec { ]) } - fn keep_on_cpu_keys(&self) -> Vec { + fn cpu_resident_tensor_keys(&self) -> Vec { vec!["image_grid_thw".to_string(), "video_grid_thw".to_string()] } } diff --git a/crates/multimodal/src/registry/qwen_vl.rs b/crates/multimodal/src/registry/qwen_vl.rs index ce5823459..26a5b47aa 100644 --- a/crates/multimodal/src/registry/qwen_vl.rs +++ b/crates/multimodal/src/registry/qwen_vl.rs @@ -88,7 +88,7 @@ impl ModelProcessorSpec for QwenVLVisionSpec { ]) } - fn keep_on_cpu_keys(&self) -> Vec { + fn cpu_resident_tensor_keys(&self) -> Vec { vec!["image_grid_thw".to_string()] } } diff --git a/crates/multimodal/src/registry/traits.rs b/crates/multimodal/src/registry/traits.rs index 6f5c6dc47..bca6bb43b 100644 --- a/crates/multimodal/src/registry/traits.rs +++ b/crates/multimodal/src/registry/traits.rs @@ -129,13 +129,11 @@ pub trait ModelProcessorSpec: Send + Sync { HashMap::from([("pixel_values".to_string(), FieldLayout::Batched)]) } - /// Tensor keys that should remain on CPU (not transferred to GPU). + /// Model-specific tensor keys that should remain CPU-resident. /// - /// In vLLM, certain model-specific tensors are marked `keep_on_cpu=True` - /// in their `MultiModalFieldConfig`. This method mirrors that per-model - /// knowledge so the router can send the hint via gRPC, avoiding the need - /// for the backend to instantiate a Python processor just to query it. - fn keep_on_cpu_keys(&self) -> Vec { + /// Backend adapters translate this model-level placement requirement to + /// their native wire or runtime representation. + fn cpu_resident_tensor_keys(&self) -> Vec { vec![] } } diff --git a/crates/multimodal/src/runtime.rs b/crates/multimodal/src/runtime.rs new file mode 100644 index 000000000..594d03b50 --- /dev/null +++ b/crates/multimodal/src/runtime.rs @@ -0,0 +1,143 @@ +#[cfg(feature = "opencv-video")] +use std::sync::{ + atomic::{AtomicUsize, Ordering}, + Arc, +}; +#[cfg(feature = "opencv-video")] +use std::time::Duration; + +use rayon::{ThreadPool, ThreadPoolBuildError, ThreadPoolBuilder}; + +const MAX_PREPROCESS_THREADS: usize = 64; + +/// Shared execution resources for multimodal decode and preprocessing. +/// +/// A runtime is owned by the application-level multimodal components and is +/// independent of any model, modality, decoder implementation, or inference +/// backend. Keeping these resources explicit avoids process-global scheduling +/// state and lets all stages share the same CPU budget. +pub struct MultimodalRuntime { + preprocess_pool: ThreadPool, + #[cfg(feature = "opencv-video")] + video_decodes: Arc, +} + +impl MultimodalRuntime { + pub fn new() -> Result { + let available_parallelism = std::thread::available_parallelism() + .map(|parallelism| parallelism.get()) + .unwrap_or(1); + let preprocess_threads = available_parallelism.min(MAX_PREPROCESS_THREADS); + let preprocess_pool = ThreadPoolBuilder::new() + .num_threads(preprocess_threads) + .thread_name(|index| format!("smg-mm-preprocess-{index}")) + .build()?; + + Ok(Self { + preprocess_pool, + #[cfg(feature = "opencv-video")] + video_decodes: Arc::new(VideoDecodeScheduler { + active: AtomicUsize::new(0), + available_parallelism, + }), + }) + } + + /// Execute CPU-heavy multimodal work in this runtime's worker pool. + pub fn run_cpu(&self, op: OP) -> R + where + OP: FnOnce() -> R + Send, + R: Send, + { + self.preprocess_pool.install(op) + } + + #[cfg(feature = "opencv-video")] + pub(crate) fn enter_video_decode(&self, coalesce_window: Duration) -> ActiveVideoDecode { + ActiveVideoDecode::enter(self.video_decodes.clone(), coalesce_window) + } +} + +#[cfg(feature = "opencv-video")] +struct VideoDecodeScheduler { + active: AtomicUsize, + available_parallelism: usize, +} + +#[cfg(feature = "opencv-video")] +pub(crate) struct ActiveVideoDecode { + scheduler: Arc, + observed: usize, +} + +#[cfg(feature = "opencv-video")] +impl ActiveVideoDecode { + fn enter(scheduler: Arc, coalesce_window: Duration) -> Self { + scheduler.active.fetch_add(1, Ordering::AcqRel); + std::thread::sleep(coalesce_window); + Self { + observed: scheduler.active.load(Ordering::Acquire), + scheduler, + } + } + + pub(crate) fn count(&self) -> usize { + self.observed + } + + pub(crate) fn available_parallelism(&self) -> usize { + self.scheduler.available_parallelism + } +} + +#[cfg(feature = "opencv-video")] +impl Drop for ActiveVideoDecode { + fn drop(&mut self) { + self.scheduler.active.fetch_sub(1, Ordering::AcqRel); + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn cpu_work_runs_in_owned_pool() { + let Ok(runtime) = MultimodalRuntime::new() else { + panic!("multimodal runtime should initialize"); + }; + + let thread_name = runtime.run_cpu(|| { + let thread = std::thread::current(); + thread.name().map(str::to_owned) + }); + + assert!(thread_name + .as_deref() + .is_some_and(|name| name.starts_with("smg-mm-preprocess-"))); + } + + #[cfg(feature = "opencv-video")] + #[test] + fn video_decode_state_is_scoped_to_runtime() { + let Ok(first_runtime) = MultimodalRuntime::new() else { + panic!("first multimodal runtime should initialize"); + }; + let Ok(second_runtime) = MultimodalRuntime::new() else { + panic!("second multimodal runtime should initialize"); + }; + + let first = first_runtime.enter_video_decode(Duration::ZERO); + let concurrent = first_runtime.enter_video_decode(Duration::ZERO); + let independent = second_runtime.enter_video_decode(Duration::ZERO); + + assert_eq!(first.count(), 1); + assert_eq!(concurrent.count(), 2); + assert_eq!(independent.count(), 1); + + drop(first); + drop(concurrent); + let next = first_runtime.enter_video_decode(Duration::ZERO); + assert_eq!(next.count(), 1); + } +} diff --git a/crates/multimodal/src/types.rs b/crates/multimodal/src/types.rs index 50aed9b01..491a4662c 100644 --- a/crates/multimodal/src/types.rs +++ b/crates/multimodal/src/types.rs @@ -1,4 +1,9 @@ -use std::{collections::HashMap, fmt, path::PathBuf, sync::Arc}; +use std::{ + collections::HashMap, + fmt, + path::PathBuf, + sync::{mpsc::Receiver, Arc, Mutex}, +}; use image::{DynamicImage, RgbImage}; use serde::{Deserialize, Serialize}; @@ -111,14 +116,78 @@ pub struct ImageFrame { /// Decoded video payload captured by the media connector. #[derive(Debug, Clone)] pub struct VideoClip { - pub frames: Vec, - pub rgb_video: Option, + frames: VideoFrames, pub raw_bytes: bytes::Bytes, pub source: VideoSource, /// Blake3 hex-digest of raw_bytes, computed at decode time. pub hash: String, } +#[derive(Debug, Clone)] +enum VideoFrames { + Dynamic(Vec), + Rgb(DecodedRgbVideo), + Stream(Arc), +} + +/// One owned sampled three-channel frame emitted by a streaming video decoder. +#[derive(Debug, Clone)] +pub struct OwnedRgbFrame { + pub width: u32, + pub height: u32, + pub data: bytes::Bytes, + pub channel_order: RgbChannelOrder, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum RgbChannelOrder { + Rgb, + Bgr, +} + +impl RgbChannelOrder { + pub(crate) fn source_channel(self, rgb_channel: usize) -> usize { + match self { + Self::Rgb => rgb_channel, + Self::Bgr => 2 - rgb_channel, + } + } +} + +/// Bounded sampled-frame stream consumed by video preprocessors. +#[derive(Debug)] +pub struct DecodedRgbFrameStream { + expected_frames: usize, + receiver: Receiver>, +} + +#[derive(Debug)] +struct DecodedRgbFrameStreamSlot { + expected_frames: usize, + stream: Mutex>, +} + +impl DecodedRgbFrameStream { + pub fn new(expected_frames: usize, receiver: Receiver>) -> Self { + Self { + expected_frames, + receiver, + } + } + + pub fn expected_frames(&self) -> usize { + self.expected_frames + } + + pub fn next_frame(&self) -> Result, String> { + match self.receiver.recv() { + Ok(Ok(frame)) => Ok(Some(frame)), + Ok(Err(error)) => Err(error), + Err(_) => Ok(None), + } + } +} + /// Borrowed RGB frame data for video preprocessors. #[derive(Debug, Clone, Copy)] pub struct RgbFrameRef<'a> { @@ -201,8 +270,7 @@ impl VideoClip { hash: String, ) -> Self { Self { - frames, - rgb_video: None, + frames: VideoFrames::Dynamic(frames), raw_bytes, source, hash, @@ -216,8 +284,24 @@ impl VideoClip { hash: String, ) -> Self { Self { - frames: Vec::new(), - rgb_video: Some(rgb_video), + frames: VideoFrames::Rgb(rgb_video), + raw_bytes, + source, + hash, + } + } + + pub fn new_rgb_stream( + stream: DecodedRgbFrameStream, + raw_bytes: bytes::Bytes, + source: VideoSource, + hash: String, + ) -> Self { + Self { + frames: VideoFrames::Stream(Arc::new(DecodedRgbFrameStreamSlot { + expected_frames: stream.expected_frames(), + stream: Mutex::new(Some(stream)), + })), raw_bytes, source, hash, @@ -225,21 +309,46 @@ impl VideoClip { } pub fn frames(&self) -> &[DynamicImage] { - &self.frames + match &self.frames { + VideoFrames::Dynamic(frames) => frames, + VideoFrames::Rgb(_) | VideoFrames::Stream(_) => &[], + } } pub fn rgb_video(&self) -> Option<&DecodedRgbVideo> { - self.rgb_video.as_ref() + match &self.frames { + VideoFrames::Rgb(video) => Some(video), + VideoFrames::Dynamic(_) | VideoFrames::Stream(_) => None, + } + } + + pub fn take_rgb_stream(&self) -> Result, String> { + let VideoFrames::Stream(stream) = &self.frames else { + return Ok(None); + }; + stream + .stream + .lock() + .map_err(|_| "decoded RGB frame stream lock is poisoned".to_string()) + .map(|mut slot| slot.take()) + } + + pub fn frame_count(&self) -> usize { + match &self.frames { + VideoFrames::Dynamic(frames) => frames.len(), + VideoFrames::Rgb(video) => video.frames.len(), + VideoFrames::Stream(stream) => stream.expected_frames, + } } pub fn materialized_frames(&self) -> Result, String> { - if !self.frames.is_empty() { - return Ok(self.frames.clone()); + match &self.frames { + VideoFrames::Dynamic(frames) => Ok(frames.clone()), + VideoFrames::Rgb(video) => video.to_dynamic_images(), + VideoFrames::Stream(_) => { + Err("streaming video frames cannot be materialized before consumption".to_string()) + } } - self.rgb_video - .as_ref() - .ok_or_else(|| "video clip has no decoded frames".to_string())? - .to_dynamic_images() } pub fn raw_bytes(&self) -> &[u8] { @@ -372,6 +481,8 @@ impl PromptReplacement { #[cfg(test)] mod tests { + use std::sync::mpsc::sync_channel; + use super::*; #[test] @@ -389,4 +500,56 @@ mod tests { let rep = PromptReplacement::repeated(Modality::Image, "", 100, 3); assert_eq!(rep.tokens, vec![100, 100, 100]); } + + #[test] + fn video_clip_representation_accessors_are_exclusive() { + let dynamic = VideoClip::new( + vec![DynamicImage::new_rgb8(2, 3)], + bytes::Bytes::new(), + VideoSource::InlineBytes, + "dynamic".to_string(), + ); + assert_eq!(dynamic.frames().len(), 1); + assert!(dynamic.rgb_video().is_none()); + assert!(dynamic.take_rgb_stream().unwrap().is_none()); + + let rgb = VideoClip::new_rgb( + DecodedRgbVideo::new( + bytes::Bytes::from_static(&[0, 0, 0]), + vec![DecodedRgbFrame { + width: 1, + height: 1, + offset: 0, + len: 3, + }], + ), + bytes::Bytes::new(), + VideoSource::InlineBytes, + "rgb".to_string(), + ); + assert!(rgb.frames().is_empty()); + assert_eq!(rgb.rgb_video().unwrap().frames.len(), 1); + assert!(rgb.take_rgb_stream().unwrap().is_none()); + + let (sender, receiver) = sync_channel(1); + sender + .send(Ok(OwnedRgbFrame { + width: 1, + height: 1, + data: bytes::Bytes::from_static(&[0, 0, 0]), + channel_order: RgbChannelOrder::Rgb, + })) + .unwrap(); + drop(sender); + let stream = VideoClip::new_rgb_stream( + DecodedRgbFrameStream::new(1, receiver), + bytes::Bytes::new(), + VideoSource::InlineBytes, + "stream".to_string(), + ); + assert!(stream.frames().is_empty()); + assert!(stream.rgb_video().is_none()); + assert!(stream.take_rgb_stream().unwrap().is_some()); + assert!(stream.take_rgb_stream().unwrap().is_none()); + } } diff --git a/crates/multimodal/src/vision/mod.rs b/crates/multimodal/src/vision/mod.rs index 02640205f..80dceeca0 100644 --- a/crates/multimodal/src/vision/mod.rs +++ b/crates/multimodal/src/vision/mod.rs @@ -37,7 +37,9 @@ pub mod transforms; // Re-export commonly used types pub use preprocessor_config::PreProcessorConfig; pub use processor::{ - ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor, VisionProcessorRegistry, + DeferredNormalizedEncoderInput, EncoderInput, ModalityPreProcessor, ModalityProcessorRegistry, + ModelSpecificValue, OutputPreference, PreprocessRequest, PreprocessedEncoderInputs, VideoInput, + VisionInput, VisionPreProcessor, VisionPreprocessRequest, VisionProcessorRegistry, }; pub use processors::{ Llama4VisionProcessor, LlavaNextProcessor, LlavaProcessor, Phi3VisionProcessor, diff --git a/crates/multimodal/src/vision/processor.rs b/crates/multimodal/src/vision/processor.rs index 000a1d505..ec51c05d6 100644 --- a/crates/multimodal/src/vision/processor.rs +++ b/crates/multimodal/src/vision/processor.rs @@ -3,14 +3,17 @@ //! This module defines the interface for model-specific vision processors //! and the common output format for preprocessed encoder inputs. -use std::{borrow::Cow, collections::HashMap}; +use std::{borrow::Cow, collections::HashMap, mem::size_of}; use anyhow::{Context, Result as AnyhowResult}; use image::DynamicImage; -use ndarray::{Array4, ArrayD}; +use ndarray::{Array4, ArrayD, IxDyn}; -use super::{preprocessor_config::PreProcessorConfig, transforms::TransformError}; -use crate::types::{FieldLayout, RgbFrameRef}; +use super::{ + preprocessor_config::PreProcessorConfig, + transforms::{par_scope, par_threads, TransformError}, +}; +use crate::types::{DecodedRgbFrameStream, FieldLayout, RgbFrameRef}; /// Helper to extract a dimension from encoder_input given an ndim-dependent axis index. /// Returns `Err` if the ndim is not 4 or 5. @@ -199,13 +202,14 @@ fn slice_1d(values: &[T], start: usize, len: usize) -> AnyhowResult<&[T]> { /// can reuse the same output contract for audio features or other encoder inputs. #[derive(Debug, Clone)] pub struct PreprocessedEncoderInputs { - /// Primary encoder input as a dynamic-dimensional float32 tensor. + /// Primary encoder input, either materialized as a dynamic-dimensional + /// float32 tensor or held in a compact representation for later assembly. /// /// For vision models this is typically the preprocessed image/video tensor. /// Shape varies by model and modality: /// - Standard: [B, C, H, W] (4D) /// - Phi3-Vision: [B, num_crops+1, C, H, W] (5D) - pub encoder_input: ArrayD, + pub encoder_input: EncoderInput, /// Number of encoder feature tokens per media item in the batch. /// @@ -231,6 +235,244 @@ pub struct PreprocessedEncoderInputs { pub model_specific: HashMap, } +/// Decoded vision payload supplied to a vision preprocessor. +pub enum VisionInput<'a> { + Images(&'a [&'a DynamicImage]), + Video(VideoInput<'a>), +} + +/// Decoded video representation supplied to a preprocessor. +pub enum VideoInput<'a> { + Frames(&'a [DynamicImage]), + Rgb(&'a [RgbFrameRef<'a>]), + RgbStream(DecodedRgbFrameStream), +} + +/// Preferred physical form of the encoder output. +/// +/// This describes a serving constraint rather than a backend or tensor dtype; +/// processors remain free to choose the most efficient compatible form. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum OutputPreference { + Materialized, + CompactAllowed, +} + +pub struct VisionPreprocessRequest<'a> { + pub input: VisionInput<'a>, + pub output: OutputPreference, +} + +/// Typed modality request supplied to the serving preprocessor contract. +/// +/// Each variant owns the matching input and configuration types. Adding audio +/// or another modality therefore requires explicit dispatch without forcing +/// it through vision-specific configuration. +pub enum PreprocessRequest<'a> { + Vision { + input: VisionInput<'a>, + output: OutputPreference, + config: &'a PreProcessorConfig, + }, +} + +/// Physical representation of a preprocessed modality encoder input. +/// +/// The variants are mutually exclusive, so callers cannot accidentally attach +/// both a dense tensor and a deferred payload or represent deferred data with an +/// empty sentinel tensor. +#[derive(Debug, Clone)] +pub enum EncoderInput { + Dense(ArrayD), + DeferredNormalized(DeferredNormalizedEncoderInput), +} + +impl EncoderInput { + pub fn shape(&self) -> &[usize] { + match self { + Self::Dense(input) => input.shape(), + Self::DeferredNormalized(input) => input.shape(), + } + } + + pub fn ndim(&self) -> usize { + self.shape().len() + } + + pub fn is_deferred(&self) -> bool { + matches!(self, Self::DeferredNormalized(_)) + } + + pub fn dense(&self) -> Result<&ArrayD, TransformError> { + match self { + Self::Dense(input) => Ok(input), + Self::DeferredNormalized(_) => Err(TransformError::ShapeError( + "encoder input must be materialized before dense access".to_string(), + )), + } + } + + pub fn deferred_normalized(&self) -> Option<&DeferredNormalizedEncoderInput> { + match self { + Self::DeferredNormalized(input) => Some(input), + Self::Dense(_) => None, + } + } + + pub fn materialize(&mut self) -> Result<(), TransformError> { + if let Self::DeferredNormalized(input) = self { + let dense = input.materialize_f32()?; + *self = Self::Dense(dense); + } + Ok(()) + } + + pub fn as_slice_memory_order(&self) -> Option<&[f32]> { + match self { + Self::Dense(input) => input.as_slice_memory_order(), + Self::DeferredNormalized(_) => None, + } + } +} + +#[derive(Debug, Clone)] +pub struct DeferredNormalizedEncoderInput { + data: Vec, + shape: Vec, + lut: Box<[[f32; 256]; 3]>, + bf16_lut: Box<[[u16; 256]; 3]>, + channel_run: usize, +} + +impl DeferredNormalizedEncoderInput { + pub fn new( + data: Vec, + shape: Vec, + lut: [[f32; 256]; 3], + channel_run: usize, + ) -> Result { + let expected = shape.iter().try_fold(1usize, |count, &dimension| { + count.checked_mul(dimension).ok_or_else(|| { + TransformError::ShapeError("deferred encoder input size overflow".to_string()) + }) + })?; + let valid_channel_layout = channel_run != 0 + && data.len().is_multiple_of(channel_run) + && (data.len() / channel_run).is_multiple_of(3); + if data.len() != expected || !valid_channel_layout { + return Err(TransformError::InvalidShape { + expected: format!( + "{expected} deferred values partitioned into complete RGB channel runs" + ), + actual: vec![data.len(), channel_run], + }); + } + let bf16_lut = lut.map(|channel| { + channel.map(|value| { + let bits = value.to_bits(); + let lsb = (bits >> 16) & 1; + (bits.wrapping_add(0x7fff + lsb) >> 16) as u16 + }) + }); + Ok(Self { + data, + shape, + lut: Box::new(lut), + bf16_lut: Box::new(bf16_lut), + channel_run, + }) + } + + pub fn len(&self) -> usize { + self.data.len() + } + + pub fn is_empty(&self) -> bool { + self.data.is_empty() + } + + pub fn shape(&self) -> &[usize] { + &self.shape + } + + pub fn materialize_f32(&self) -> Result, TransformError> { + let mut values = vec![0.0; self.data.len()]; + self.fill_f32(&mut values)?; + ArrayD::from_shape_vec(IxDyn(&self.shape), values).map_err(|error| { + TransformError::ShapeError(format!( + "failed to materialize deferred encoder input: {error}" + )) + }) + } + + pub fn fill_f32(&self, output: &mut [f32]) -> Result<(), TransformError> { + if output.len() != self.data.len() { + return Err(TransformError::InvalidShape { + expected: format!("{} normalized encoder values", self.data.len()), + actual: vec![output.len()], + }); + } + self.fill_parallel(output, size_of::(), |group, input, output| { + let lut = &self.lut[group % 3]; + for (&value, output) in input.iter().zip(output) { + *output = lut[value as usize]; + } + }); + Ok(()) + } + + pub fn fill_bf16_le_bytes(&self, output: &mut [u8]) -> Result<(), TransformError> { + if output.len() != self.data.len() * size_of::() { + return Err(TransformError::InvalidShape { + expected: format!("{} BF16 encoder bytes", self.data.len() * 2), + actual: vec![output.len()], + }); + } + self.fill_parallel(output, size_of::(), |group, input, output| { + let lut = &self.bf16_lut[group % 3]; + for (&value, output) in input.iter().zip(output.chunks_exact_mut(2)) { + output.copy_from_slice(&lut[value as usize].to_le_bytes()); + } + }); + Ok(()) + } + + fn fill_parallel(&self, output: &mut [T], output_bytes_per_value: usize, fill: F) + where + T: Send, + F: Fn(usize, &[u8], &mut [T]) + Copy + Send + Sync, + { + let groups = self.data.len().div_ceil(self.channel_run); + let workers = par_threads(self.data.len() * output_bytes_per_value, groups); + let groups_per_task = groups.div_ceil(workers); + let input_values_per_task = groups_per_task * self.channel_run; + let output_values_per_task = + input_values_per_task * output_bytes_per_value / size_of::(); + par_scope(|scope| { + for (task, (input, output)) in self + .data + .chunks(input_values_per_task) + .zip(output.chunks_mut(output_values_per_task)) + .enumerate() + { + let first_group = task * groups_per_task; + scope.spawn(move |_| { + for (group_offset, (input, output)) in + input + .chunks(self.channel_run) + .zip(output.chunks_mut( + self.channel_run * output_bytes_per_value / size_of::(), + )) + .enumerate() + { + fill(first_group + group_offset, input, output); + } + }); + } + }); + } +} + impl PreprocessedEncoderInputs { /// Create a new PreprocessedEncoderInputs with required fields (4D encoder input). pub fn new( @@ -239,7 +481,7 @@ impl PreprocessedEncoderInputs { item_sizes: Vec<(u32, u32)>, ) -> Self { Self { - encoder_input: encoder_input.into_dyn(), + encoder_input: EncoderInput::Dense(encoder_input.into_dyn()), feature_token_counts, item_sizes, model_specific: HashMap::new(), @@ -255,19 +497,45 @@ impl PreprocessedEncoderInputs { item_sizes: Vec<(u32, u32)>, ) -> Self { Self { - encoder_input, + encoder_input: EncoderInput::Dense(encoder_input), feature_token_counts, item_sizes, model_specific: HashMap::new(), } } + pub fn new_deferred_normalized( + deferred_encoder_input: DeferredNormalizedEncoderInput, + feature_token_counts: Vec, + item_sizes: Vec<(u32, u32)>, + ) -> Self { + Self { + encoder_input: EncoderInput::DeferredNormalized(deferred_encoder_input), + feature_token_counts, + item_sizes, + model_specific: HashMap::new(), + } + } + + pub fn materialize_encoder_input(&mut self) -> Result<(), TransformError> { + self.encoder_input.materialize() + } + /// Add a model-specific value. pub fn with_extra(mut self, key: impl Into, value: ModelSpecificValue) -> Self { self.model_specific.insert(key.into(), value); self } + /// Attach the complete model-specific output set produced by a processor. + pub fn with_model_specific( + mut self, + model_specific: HashMap, + ) -> Self { + self.model_specific = model_specific; + self + } + /// Get the number of media items represented by this preprocessed batch. pub fn batch_size(&self) -> usize { self.item_sizes.len() @@ -317,11 +585,12 @@ impl PreprocessedEncoderInputs { } /// Get the primary encoder input as a flat f32 slice without copying if possible. - pub fn encoder_input_flat(&self) -> Cow<'_, [f32]> { - match self.encoder_input.as_slice() { + pub fn encoder_input_flat(&self) -> Result, TransformError> { + let encoder_input = self.encoder_input.dense()?; + Ok(match encoder_input.as_slice() { Some(slice) => Cow::Borrowed(slice), - None => Cow::Owned(self.encoder_input.iter().copied().collect()), - } + None => Cow::Owned(encoder_input.iter().copied().collect()), + }) } /// Get the shape of the primary encoder input as a vector. @@ -357,6 +626,39 @@ impl PreprocessedEncoderInputs { } } +/// Modality-neutral preprocessing contract used by serving pipelines. +pub trait ModalityPreProcessor: Send + Sync { + /// Stable processor family name for diagnostics and registry inspection. + fn processor_name(&self) -> &'static str; + + fn preprocess_input( + &self, + request: PreprocessRequest<'_>, + ) -> Result; +} + +impl ModalityPreProcessor for T +where + T: VisionPreProcessor + ?Sized, +{ + fn processor_name(&self) -> &'static str { + self.model_name() + } + + fn preprocess_input( + &self, + request: PreprocessRequest<'_>, + ) -> Result { + match request { + PreprocessRequest::Vision { + input, + output, + config, + } => self.preprocess_vision_input(VisionPreprocessRequest { input, output }, config), + } + } +} + /// Trait for model-specific vision preprocessors. /// /// Each vision model (LLaVA, Qwen-VL, Phi3-Vision, etc.) implements this trait @@ -382,34 +684,27 @@ pub trait VisionPreProcessor: Send + Sync { config: &PreProcessorConfig, ) -> Result; - /// Preprocess one decoded video clip represented as sampled frames. - /// - /// Implementations that support video should emit the same primary - /// `encoder_input` tensor shape used by the image path, plus video-specific - /// model metadata such as `video_grid_thw`. - fn preprocess_video( + /// Preprocess a modality payload independently of its decoded + /// representation. Implementations may select a compact output only when + /// the request permits it. + fn preprocess_vision_input( &self, - _frames: &[DynamicImage], - _config: &PreProcessorConfig, - ) -> Result { - Err(TransformError::ShapeError(format!( - "{} does not support video preprocessing", - self.model_name() - ))) - } - - /// Preprocess one decoded video clip represented as borrowed RGB frame - /// buffers. Implementations can override this to avoid materializing - /// `DynamicImage` objects after media decode. - fn preprocess_video_rgb( - &self, - _frames: &[RgbFrameRef<'_>], - _config: &PreProcessorConfig, + request: VisionPreprocessRequest<'_>, + config: &PreProcessorConfig, ) -> Result { - Err(TransformError::ShapeError(format!( - "{} does not support RGB video preprocessing", - self.model_name() - ))) + match request.input { + VisionInput::Images(images) => { + let owned = images + .iter() + .map(|image| (**image).clone()) + .collect::>(); + self.preprocess(&owned, config) + } + VisionInput::Video(_) => Err(TransformError::ShapeError(format!( + "{} does not support video preprocessing", + self.model_name() + ))), + } } /// Calculate the number of vision tokens for a given image size. @@ -434,12 +729,12 @@ pub trait VisionPreProcessor: Send + Sync { } } -/// Registry of available vision processors. -pub struct VisionProcessorRegistry { - processors: HashMap>, +/// Registry of model-specific modality preprocessors. +pub struct ModalityProcessorRegistry { + processors: HashMap>, } -impl VisionProcessorRegistry { +impl ModalityProcessorRegistry { /// Create a new empty registry. pub fn new() -> Self { Self { @@ -448,8 +743,12 @@ impl VisionProcessorRegistry { } /// Register a processor for a model pattern. - pub fn register(&mut self, pattern: impl Into, processor: Box) { - self.processors.insert(pattern.into(), processor); + pub fn register

(&mut self, pattern: impl Into, processor: P) + where + P: ModalityPreProcessor + 'static, + { + self.processors + .insert(pattern.into().to_lowercase(), Box::new(processor)); } /// Find a processor for the given model ID, falling back to model_type. @@ -459,19 +758,20 @@ impl VisionProcessorRegistry { &self, model_id: &str, model_type: Option<&str>, - ) -> Option<&dyn VisionPreProcessor> { + ) -> Option<&dyn ModalityPreProcessor> { self.find_in_candidate(model_id) .or_else(|| model_type.and_then(|mt| self.find_in_candidate(mt))) } - fn find_in_candidate(&self, candidate: &str) -> Option<&dyn VisionPreProcessor> { + fn find_in_candidate(&self, candidate: &str) -> Option<&dyn ModalityPreProcessor> { let candidate = candidate.to_lowercase(); - for (pattern, processor) in &self.processors { - if candidate.contains(&pattern.to_lowercase()) { - return Some(processor.as_ref()); - } - } - None + self.processors + .iter() + .filter(|(pattern, _)| candidate.contains(pattern.as_str())) + .max_by(|(left, _), (right, _)| { + left.len().cmp(&right.len()).then_with(|| left.cmp(right)) + }) + .map(|(_, processor)| processor.as_ref()) } /// Get list of supported model patterns. @@ -480,13 +780,16 @@ impl VisionProcessorRegistry { } } -impl Default for VisionProcessorRegistry { +/// Backward-compatible name for callers that only register vision processors. +pub type VisionProcessorRegistry = ModalityProcessorRegistry; + +impl Default for ModalityProcessorRegistry { fn default() -> Self { Self::new() } } -impl VisionProcessorRegistry { +impl ModalityProcessorRegistry { /// Create a registry with all built-in processors registered. /// /// Currently registers: @@ -501,116 +804,50 @@ impl VisionProcessorRegistry { let mut registry = Self::new(); // LLaVA-NeXT (v1.6+, anyres multi-crop) - registry.register( - "llava-next", - Box::new(super::processors::LlavaNextProcessor::new()), - ); - registry.register( - "llava_next", - Box::new(super::processors::LlavaNextProcessor::new()), - ); - registry.register( - "llava-v1.6", - Box::new(super::processors::LlavaNextProcessor::new()), - ); + registry.register("llava-next", super::processors::LlavaNextProcessor::new()); + registry.register("llava_next", super::processors::LlavaNextProcessor::new()); + registry.register("llava-v1.6", super::processors::LlavaNextProcessor::new()); // Standard LLaVA (v1.5, single-patch). // Use specific patterns so they don't accidentally match LLaVA-NeXT // model IDs like "llava-v1.6-*". - registry.register( - "llava-1.5", - Box::new(super::processors::LlavaProcessor::new()), - ); - registry.register( - "llava-v1.5", - Box::new(super::processors::LlavaProcessor::new()), - ); + registry.register("llava-1.5", super::processors::LlavaProcessor::new()); + registry.register("llava-v1.5", super::processors::LlavaProcessor::new()); // Register Qwen3-VL first (more specific pattern - must match before qwen2) - registry.register( - "qwen3-vl", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - registry.register( - "qwen3_vl", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); + registry.register("qwen3-vl", super::processors::Qwen3VLProcessor::new()); + registry.register("qwen3_vl", super::processors::Qwen3VLProcessor::new()); // Qwen3.5 family (and Qwen3.6: same arch) reuses Qwen3-VL preprocessing. - registry.register( - "qwen3.5", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - registry.register( - "qwen3_5", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - registry.register( - "qwen3.6", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); - registry.register( - "qwen3_6", - Box::new(super::processors::Qwen3VLProcessor::new()), - ); + registry.register("qwen3.5", super::processors::Qwen3VLProcessor::new()); + registry.register("qwen3_5", super::processors::Qwen3VLProcessor::new()); + registry.register("qwen3.6", super::processors::Qwen3VLProcessor::new()); + registry.register("qwen3_6", super::processors::Qwen3VLProcessor::new()); // Register Qwen2-VL (matches Qwen/Qwen2-VL-*, etc.) - registry.register( - "qwen2-vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); - registry.register( - "qwen2_vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); + registry.register("qwen2-vl", super::processors::Qwen2VLProcessor::new()); + registry.register("qwen2_vl", super::processors::Qwen2VLProcessor::new()); // Register Qwen2.5-VL (uses identical preprocessing to Qwen2-VL) - registry.register( - "qwen2.5-vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); - registry.register( - "qwen2_5-vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); - registry.register( - "qwen2_5_vl", - Box::new(super::processors::Qwen2VLProcessor::new()), - ); + registry.register("qwen2.5-vl", super::processors::Qwen2VLProcessor::new()); + registry.register("qwen2_5-vl", super::processors::Qwen2VLProcessor::new()); + registry.register("qwen2_5_vl", super::processors::Qwen2VLProcessor::new()); // Register Phi3-Vision registry.register( "phi-3-vision", - Box::new(super::processors::Phi3VisionProcessor::new()), - ); - registry.register( - "phi3-vision", - Box::new(super::processors::Phi3VisionProcessor::new()), - ); - registry.register( - "phi3_v", - Box::new(super::processors::Phi3VisionProcessor::new()), + super::processors::Phi3VisionProcessor::new(), ); + registry.register("phi3-vision", super::processors::Phi3VisionProcessor::new()); + registry.register("phi3_v", super::processors::Phi3VisionProcessor::new()); // Register LLaMA 4 Vision - registry.register( - "llama-4", - Box::new(super::processors::Llama4VisionProcessor::new()), - ); - registry.register( - "llama4", - Box::new(super::processors::Llama4VisionProcessor::new()), - ); + registry.register("llama-4", super::processors::Llama4VisionProcessor::new()); + registry.register("llama4", super::processors::Llama4VisionProcessor::new()); // Register Kimi-K2.5 Vision - registry.register( - "kimi-k2", - Box::new(super::processors::KimiK25Processor::new()), - ); - registry.register( - "kimi_k2", - Box::new(super::processors::KimiK25Processor::new()), - ); + registry.register("kimi-k2", super::processors::KimiK25Processor::new()); + registry.register("kimi_k2", super::processors::KimiK25Processor::new()); registry } @@ -623,6 +860,23 @@ mod tests { use super::*; use crate::vision::processors::LlavaProcessor; + struct NonVisionProcessor; + + impl ModalityPreProcessor for NonVisionProcessor { + fn processor_name(&self) -> &'static str { + "non-vision" + } + + fn preprocess_input( + &self, + _request: PreprocessRequest<'_>, + ) -> Result { + Err(TransformError::ShapeError( + "test processor has no payload implementation".to_string(), + )) + } + } + #[test] fn test_preprocessed_encoder_inputs_accessors() { let encoder_input = Array4::::zeros((2, 3, 336, 336)); @@ -701,14 +955,42 @@ mod tests { encoder_input[[0, 0, 1, 1]] = 4.0; let inputs = PreprocessedEncoderInputs::new(encoder_input, vec![4], vec![(2, 2)]); - let flat = inputs.encoder_input_flat(); + let flat = inputs.encoder_input_flat().unwrap(); assert_eq!(flat, vec![1.0, 2.0, 3.0, 4.0]); } + #[test] + fn deferred_encoder_input_validates_rgb_channel_runs() { + let lut = [[0.0; 256]; 3]; + + assert!(DeferredNormalizedEncoderInput::new(vec![0; 6], vec![2, 3], lut, 4).is_err()); + assert!(DeferredNormalizedEncoderInput::new(vec![0; 4], vec![1, 4], lut, 1).is_err()); + } + + #[test] + fn deferred_encoder_input_accessors_use_logical_shape() { + let lut = std::array::from_fn(|channel| { + std::array::from_fn(|value| value as f32 + channel as f32) + }); + let deferred = + DeferredNormalizedEncoderInput::new(vec![0; 12], vec![1, 3, 2, 2], lut, 4).unwrap(); + let mut inputs = + PreprocessedEncoderInputs::new_deferred_normalized(deferred, vec![1], vec![(2, 2)]); + + assert_eq!(inputs.channels().unwrap(), 3); + assert_eq!(inputs.height().unwrap(), 2); + assert_eq!(inputs.width().unwrap(), 2); + assert_eq!(inputs.encoder_input_shape(), vec![1, 3, 2, 2]); + assert!(inputs.encoder_input_flat().is_err()); + + inputs.materialize_encoder_input().unwrap(); + assert_eq!(inputs.encoder_input_flat().unwrap().len(), 12); + } + #[test] fn test_registry_with_defaults() { - let registry = VisionProcessorRegistry::with_defaults(); + let registry = ModalityProcessorRegistry::with_defaults(); // Should find LLaVA processor assert!(registry.find("llava-hf/llava-1.5-7b-hf", None).is_some()); @@ -724,50 +1006,69 @@ mod tests { // Get the processor and check model name let processor = registry.find("llava-hf/llava-1.5-7b-hf", None).unwrap(); - assert_eq!(processor.model_name(), "llava"); + assert_eq!(processor.processor_name(), "llava"); } #[test] fn test_registry_find() { - let mut registry = VisionProcessorRegistry::new(); + let mut registry = ModalityProcessorRegistry::new(); // Create a mock processor using LlavaProcessor - registry.register("test-model", Box::new(LlavaProcessor::new())); + registry.register("test-model", LlavaProcessor::new()); assert!(registry.find("test-model-7b", None).is_some()); assert!(registry.find("TEST-MODEL", None).is_some()); assert!(registry.find("other-model", None).is_none()); } + #[test] + fn test_registry_accepts_non_vision_processor() { + let mut registry = ModalityProcessorRegistry::new(); + registry.register("audio-model", NonVisionProcessor); + + let processor = registry.find("vendor/audio-model", None).unwrap(); + assert_eq!(processor.processor_name(), "non-vision"); + } + + #[test] + fn test_registry_prefers_most_specific_pattern_deterministically() { + let mut registry = ModalityProcessorRegistry::new(); + registry.register("audio", NonVisionProcessor); + registry.register("audio-special", LlavaProcessor::new()); + + let processor = registry.find("vendor/audio-special-model", None).unwrap(); + assert_eq!(processor.processor_name(), "llava"); + } + #[test] fn test_registry_find_falls_back_to_model_type() { - let registry = VisionProcessorRegistry::with_defaults(); + let registry = ModalityProcessorRegistry::with_defaults(); assert!(registry.find("custom-model", None).is_none()); let processor = registry .find("custom-model", Some("qwen3_vl")) .expect("qwen3 processor by model_type"); - assert_eq!(processor.model_name(), "qwen3-vl"); + assert_eq!(processor.processor_name(), "qwen3-vl"); } #[test] fn test_registry_find_preserves_fast_path() { - let registry = VisionProcessorRegistry::with_defaults(); + let registry = ModalityProcessorRegistry::with_defaults(); let processor = registry .find("Qwen3-VL-30B-A3B-Instruct", Some("qwen2_vl")) .expect("qwen3 processor by model_id"); - assert_eq!(processor.model_name(), "qwen3-vl"); + assert_eq!(processor.processor_name(), "qwen3-vl"); } #[test] fn test_registry_find_phi3_model_type_fallback() { - let registry = VisionProcessorRegistry::with_defaults(); + let registry = ModalityProcessorRegistry::with_defaults(); let processor = registry .find("custom-model", Some("phi3_v")) .expect("phi3 processor by model_type"); - assert_eq!(processor.model_name(), "phi3-vision"); + assert_eq!(processor.processor_name(), "phi3-vision"); } } diff --git a/crates/multimodal/src/vision/processors/kimi_k25.rs b/crates/multimodal/src/vision/processors/kimi_k25.rs index 16f679851..3a05fe577 100644 --- a/crates/multimodal/src/vision/processors/kimi_k25.rs +++ b/crates/multimodal/src/vision/processors/kimi_k25.rs @@ -496,7 +496,7 @@ mod tests { let image = create_test_image(100, 100, Rgb([255, 255, 255])); let result = p.preprocess(&[image], &config).unwrap(); - let flat = result.encoder_input_flat(); + let flat = result.encoder_input_flat().unwrap(); // Padded region should be normalized black (-1.0) let has_neg_ones = flat.iter().any(|&v| (v - (-1.0)).abs() < 1e-6); assert!( diff --git a/crates/multimodal/src/vision/processors/llama4_vision.rs b/crates/multimodal/src/vision/processors/llama4_vision.rs index 033adc128..5722fc81a 100644 --- a/crates/multimodal/src/vision/processors/llama4_vision.rs +++ b/crates/multimodal/src/vision/processors/llama4_vision.rs @@ -501,12 +501,12 @@ impl VisionPreProcessor for Llama4VisionProcessor { ModelSpecificValue::int_1d(patches_per_image), ); - Ok(PreprocessedEncoderInputs { - encoder_input: encoder_input.into_dyn(), + Ok(PreprocessedEncoderInputs::new_dynamic( + encoder_input.into_dyn(), feature_token_counts, item_sizes, - model_specific, - }) + ) + .with_model_specific(model_specific)) } fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { @@ -627,7 +627,7 @@ mod tests { assert!(result.feature_token_counts[0] > 0); // Check pixel values are normalized - let flat = result.encoder_input_flat(); + let flat = result.encoder_input_flat().unwrap(); assert!(flat.iter().all(|&v| (-1.5..=1.5).contains(&v))); } diff --git a/crates/multimodal/src/vision/processors/phi3_vision.rs b/crates/multimodal/src/vision/processors/phi3_vision.rs index a266d4940..a2f48abd0 100644 --- a/crates/multimodal/src/vision/processors/phi3_vision.rs +++ b/crates/multimodal/src/vision/processors/phi3_vision.rs @@ -373,12 +373,10 @@ impl VisionPreProcessor for Phi3VisionProcessor { actual: shape.clone(), })?; - Ok(PreprocessedEncoderInputs { - encoder_input, - feature_token_counts: all_num_tokens, - item_sizes: all_image_sizes, - model_specific, - }) + Ok( + PreprocessedEncoderInputs::new_dynamic(encoder_input, all_num_tokens, all_image_sizes) + .with_model_specific(model_specific), + ) } fn calculate_num_tokens(&self, width: u32, height: u32, _config: &PreProcessorConfig) -> usize { diff --git a/crates/multimodal/src/vision/processors/phi4_vision.rs b/crates/multimodal/src/vision/processors/phi4_vision.rs index eb74d7c10..38a20f312 100644 --- a/crates/multimodal/src/vision/processors/phi4_vision.rs +++ b/crates/multimodal/src/vision/processors/phi4_vision.rs @@ -543,12 +543,12 @@ impl VisionPreProcessor for Phi4VisionProcessor { }, ); - Ok(PreprocessedEncoderInputs { - encoder_input: encoder_input.into_dyn(), + Ok(PreprocessedEncoderInputs::new_dynamic( + encoder_input.into_dyn(), feature_token_counts, item_sizes, - model_specific, - }) + ) + .with_model_specific(model_specific)) } fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { @@ -675,7 +675,7 @@ mod tests { assert!(result.feature_token_counts[0] > 256); // At least global tokens // Check pixel values are normalized - let flat = result.encoder_input_flat(); + let flat = result.encoder_input_flat().unwrap(); assert!(flat.iter().all(|&v| (-1.5..=1.5).contains(&v))); } diff --git a/crates/multimodal/src/vision/processors/pixtral.rs b/crates/multimodal/src/vision/processors/pixtral.rs index c414f5d91..29e6ab7ab 100644 --- a/crates/multimodal/src/vision/processors/pixtral.rs +++ b/crates/multimodal/src/vision/processors/pixtral.rs @@ -248,12 +248,12 @@ impl VisionPreProcessor for PixtralProcessor { }, ); - Ok(PreprocessedEncoderInputs { - encoder_input: batch_tensor, + Ok(PreprocessedEncoderInputs::new_dynamic( + batch_tensor, feature_token_counts, - item_sizes: original_sizes, - model_specific, - }) + original_sizes, + ) + .with_model_specific(model_specific)) } fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { diff --git a/crates/multimodal/src/vision/processors/qwen2_vl.rs b/crates/multimodal/src/vision/processors/qwen2_vl.rs index 6e2cfff80..d95f95815 100644 --- a/crates/multimodal/src/vision/processors/qwen2_vl.rs +++ b/crates/multimodal/src/vision/processors/qwen2_vl.rs @@ -24,7 +24,9 @@ use image::DynamicImage; use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase}; use crate::vision::{ preprocessor_config::PreProcessorConfig, - processor::{PreprocessedEncoderInputs, VisionPreProcessor}, + processor::{ + PreprocessedEncoderInputs, VisionInput, VisionPreProcessor, VisionPreprocessRequest, + }, transforms::TransformError, }; @@ -231,6 +233,21 @@ impl VisionPreProcessor for Qwen2VLProcessor { processor.inner.preprocess(images, config) } + fn preprocess_vision_input( + &self, + request: VisionPreprocessRequest<'_>, + config: &PreProcessorConfig, + ) -> Result { + if matches!(&request.input, VisionInput::Video(_)) { + return Err(TransformError::ShapeError(format!( + "{} does not support video preprocessing", + self.model_name() + ))); + } + let processor = self.with_preprocessor_config(config); + processor.inner.preprocess_vision_input(request, config) + } + fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { let processor = self.with_preprocessor_config(config); processor.inner.calculate_num_tokens(width, height, config) @@ -389,7 +406,7 @@ mod tests { assert!(result.encoder_input.shape()[0] > 0); // total_patches > 0 // Check pixel values are normalized - let flat = result.encoder_input_flat(); + let flat = result.encoder_input_flat().unwrap(); // After normalization with CLIP mean/std, gray (0.5) should be near 0 // (0.5 - 0.48) / 0.27 ≈ 0.07 assert!(flat.iter().all(|&v| v.abs() < 1.0)); // Should be normalized diff --git a/crates/multimodal/src/vision/processors/qwen3_vl.rs b/crates/multimodal/src/vision/processors/qwen3_vl.rs index e163967b4..f1a019be9 100644 --- a/crates/multimodal/src/vision/processors/qwen3_vl.rs +++ b/crates/multimodal/src/vision/processors/qwen3_vl.rs @@ -21,13 +21,10 @@ use std::ops::Deref; use image::DynamicImage; use super::qwen_vl_base::{QwenVLConfig, QwenVLProcessorBase}; -use crate::{ - types::RgbFrameRef, - vision::{ - preprocessor_config::PreProcessorConfig, - processor::{PreprocessedEncoderInputs, VisionPreProcessor}, - transforms::TransformError, - }, +use crate::vision::{ + preprocessor_config::PreProcessorConfig, + processor::{PreprocessedEncoderInputs, VisionPreProcessor, VisionPreprocessRequest}, + transforms::TransformError, }; /// Qwen3-VL normalization mean values (simple [0.5, 0.5, 0.5]). @@ -239,22 +236,13 @@ impl VisionPreProcessor for Qwen3VLProcessor { processor.inner.preprocess(images, config) } - fn preprocess_video( + fn preprocess_vision_input( &self, - frames: &[DynamicImage], + request: VisionPreprocessRequest<'_>, config: &PreProcessorConfig, ) -> Result { let processor = self.with_preprocessor_config(config); - processor.inner.preprocess_video(frames, config) - } - - fn preprocess_video_rgb( - &self, - frames: &[RgbFrameRef<'_>], - config: &PreProcessorConfig, - ) -> Result { - let processor = self.with_preprocessor_config(config); - processor.inner.preprocess_video_rgb(frames, config) + processor.inner.preprocess_vision_input(request, config) } fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { @@ -278,7 +266,16 @@ mod tests { use image::{Rgb, RgbImage}; use super::*; - use crate::vision::{preprocessor_config::PatchSize, processor::ModelSpecificValue}; + use crate::{ + vision::{ + preprocessor_config::PatchSize, + processor::{ + ModalityPreProcessor, ModelSpecificValue, OutputPreference, PreprocessRequest, + VideoInput, VisionInput, + }, + }, + RgbFrameRef, + }; fn create_test_image(width: u32, height: u32, color: Rgb) -> DynamicImage { DynamicImage::from(RgbImage::from_pixel(width, height, color)) @@ -391,7 +388,7 @@ mod tests { assert!(result.encoder_input.shape()[0] > 0); // total_patches > 0 // Check pixel values are normalized - let flat = result.encoder_input_flat(); + let flat = result.encoder_input_flat().unwrap(); // After normalization with [0.5, 0.5, 0.5] mean/std: // (0.5 - 0.5) / 0.5 = 0.0 for gray // Values should be in [-1, 1] range @@ -467,7 +464,13 @@ mod tests { create_test_image(640, 480, Rgb([200, 200, 200])), ]; - let result = processor.preprocess_video(&frames, &config).unwrap(); + let result = processor + .preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Video(VideoInput::Frames(&frames)), + output: OutputPreference::Materialized, + config: &config, + }) + .unwrap(); assert_eq!(result.encoder_input.ndim(), 2); assert_eq!(result.feature_token_counts.len(), 1); assert!(result.model_specific.contains_key("video_grid_thw")); @@ -512,9 +515,19 @@ mod tests { }) .collect(); - let dynamic = processor.preprocess_video(&frames, &config).unwrap(); + let dynamic = processor + .preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Video(VideoInput::Frames(&frames)), + output: OutputPreference::Materialized, + config: &config, + }) + .unwrap(); let rgb = processor - .preprocess_video_rgb(&rgb_frames, &config) + .preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Video(VideoInput::Rgb(&rgb_frames)), + output: OutputPreference::Materialized, + config: &config, + }) .unwrap(); assert_eq!(rgb.encoder_input.shape(), dynamic.encoder_input.shape()); diff --git a/crates/multimodal/src/vision/processors/qwen_vl_base.rs b/crates/multimodal/src/vision/processors/qwen_vl_base.rs index 18a7639ab..339b55405 100644 --- a/crates/multimodal/src/vision/processors/qwen_vl_base.rs +++ b/crates/multimodal/src/vision/processors/qwen_vl_base.rs @@ -27,13 +27,17 @@ use image::{imageops::FilterType, DynamicImage, GenericImageView}; use ndarray::{Array2, Array3}; use crate::{ - types::RgbFrameRef, + types::{DecodedRgbFrameStream, OwnedRgbFrame, RgbChannelOrder, RgbFrameRef}, vision::{ preprocessor_config::PreProcessorConfig, - processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor}, + processor::{ + DeferredNormalizedEncoderInput, ModelSpecificValue, OutputPreference, + PreprocessedEncoderInputs, VideoInput, VisionInput, VisionPreProcessor, + VisionPreprocessRequest, + }, transforms::{ - par_threads, pil_to_filter, resize, resize_bicubic_pil, resize_bicubic_pil_rgb, - resize_rgb_bytes, rgb_bytes, to_tensor, to_tensor_and_normalize, TransformError, + par_scope, par_threads, pil_to_filter, resize, resize_bicubic_pil, + resize_bicubic_pil_rgb, resize_rgb_bytes, rgb_bytes, PilBicubicRgbPlan, TransformError, }, }, }; @@ -83,6 +87,100 @@ struct VideoFrameRgb<'a> { data: Cow<'a, [u8]>, } +struct QwenImagePlan { + target_width: u32, + target_height: u32, + needs_resize: bool, + grid_t: usize, + grid_h: usize, + grid_w: usize, + num_patches: usize, + patch_values: usize, + tokens: usize, +} + +fn validate_streamed_qwen_frame( + frame: &OwnedRgbFrame, + expected_width: u32, + expected_height: u32, + expected_channel_order: RgbChannelOrder, +) -> Result<(), TransformError> { + if frame.width != expected_width || frame.height != expected_height { + return Err(TransformError::InvalidShape { + expected: format!("uniform RGB frames of {expected_width}x{expected_height}"), + actual: vec![frame.width as usize, frame.height as usize], + }); + } + let expected_len = (frame.width as usize) + .checked_mul(frame.height as usize) + .and_then(|pixels| pixels.checked_mul(3)) + .ok_or_else(|| { + TransformError::ShapeError(format!( + "video frame dimensions are too large: {}x{}", + frame.width, frame.height + )) + })?; + if frame.data.len() != expected_len { + return Err(TransformError::InvalidShape { + expected: format!( + "RGB frame byte length {expected_len} for {}x{}", + frame.width, frame.height + ), + actual: vec![frame.data.len()], + }); + } + if frame.channel_order != expected_channel_order { + return Err(TransformError::ShapeError( + "decoded video stream changed channel order between frames".to_string(), + )); + } + Ok(()) +} + +fn resize_rgb_frame_to_raw( + frame: RgbFrameRef<'_>, + target_width: u32, + target_height: u32, + filter: FilterType, +) -> Result, TransformError> { + // BICUBIC (Qwen default) uses the PIL-compatible path, same as the image + // path; other filters keep the SIMD resizer. + let resized = if filter == FilterType::CatmullRom { + resize_bicubic_pil_rgb( + frame.data, + frame.width, + frame.height, + target_width, + target_height, + )? + } else { + resize_rgb_bytes( + frame.data, + frame.width, + frame.height, + target_width, + target_height, + filter, + )? + }; + Ok(resized.into_raw()) +} + +fn resize_dynamic_frame_to_raw( + frame: &DynamicImage, + target_width: u32, + target_height: u32, + filter: FilterType, +) -> (usize, usize, Vec) { + let resized = if filter == FilterType::CatmullRom { + resize_bicubic_pil(frame, target_width, target_height) + } else { + resize(frame, target_width, target_height, filter) + }; + let (width, height, data) = rgb_bytes(&resized); + (width, height, data.into_owned()) +} + /// Generic Qwen VL image processor. /// /// This struct implements the shared preprocessing logic for all Qwen VL @@ -203,6 +301,27 @@ impl QwenVLProcessorBase { Ok((h_bar, w_bar)) } + fn validate_unresized_patch_dimensions( + &self, + height: usize, + width: usize, + ) -> Result<(), TransformError> { + let factor = self.get_factor(); + if height == 0 || width == 0 { + return Err(TransformError::InvalidShape { + expected: "non-zero dimensions".to_string(), + actual: vec![height, width], + }); + } + if !height.is_multiple_of(factor) || !width.is_multiple_of(factor) { + return Err(TransformError::InvalidShape { + expected: format!("height and width divisible by factor ({factor})"), + actual: vec![height, width], + }); + } + Ok(()) + } + /// Smart resize for Qwen3-style video processors. /// /// Unlike image resize, the pixel budget is applied to the full sampled @@ -367,7 +486,7 @@ impl QwenVLProcessorBase { } else { let chunk_blocks = n_blocks.div_ceil(nthreads); let planes_ref = &planes; - std::thread::scope(|s| { + par_scope(|s| { let mut rest = &mut *region; let mut b0 = 0usize; while b0 < n_blocks { @@ -375,7 +494,7 @@ impl QwenVLProcessorBase { let (band, tail) = rest.split_at_mut(nb * block_out); rest = tail; let start = b0; - s.spawn(move || { + s.spawn(move |_| { Self::patchify_block_band( planes_ref, width, @@ -559,177 +678,794 @@ impl QwenVLProcessorBase { } let merged_patch = merge_size * patch_size; - for pr in 0..grid_h / merge_size { - for pc in 0..grid_w / merge_size { - let y0 = pr * merged_patch; - let x0 = pc * merged_patch; - - for mh in 0..merge_size { - for mw in 0..merge_size { - for (c, lut_c) in lut.iter().enumerate().take(3) { - for frame in frames { - let raw = frame.data.as_ref(); - for py in 0..patch_size { - let row = - (y0 + mh * patch_size + py) * width + x0 + mw * patch_size; - let mut src_idx = row * 3 + c; - let dst_end = *out_idx + patch_size; - for dst in &mut output[*out_idx..dst_end] { - *dst = lut_c[raw[src_idx] as usize]; - src_idx += 3; - } - *out_idx = dst_end; - } - } - } - } + let pr_blocks = grid_h / merge_size; + let pc_blocks = grid_w / merge_size; + let n_blocks = pr_blocks * pc_blocks; + let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; + let base_idx = *out_idx; + let patch_values = n_blocks.checked_mul(block_out).ok_or_else(|| { + TransformError::ShapeError("Qwen video patch output size overflow".to_string()) + })?; + let end_idx = base_idx.checked_add(patch_values).ok_or_else(|| { + TransformError::ShapeError("Qwen video patch output range overflow".to_string()) + })?; + let region = output.get_mut(base_idx..end_idx).ok_or_else(|| { + TransformError::ShapeError("Qwen video patch output range out of bounds".to_string()) + })?; + let nthreads = par_threads(region.len() * 4, n_blocks); + if nthreads <= 1 { + Self::patchify_video_rgb_block_band( + frames, + width, + patch_size, + merge_size, + merged_patch, + pc_blocks, + 0, + region, + lut, + ); + } else { + let chunk_blocks = n_blocks.div_ceil(nthreads); + par_scope(|s| { + let mut rest = &mut *region; + let mut b0 = 0usize; + while b0 < n_blocks { + let nb = chunk_blocks.min(n_blocks - b0); + let (band, tail) = rest.split_at_mut(nb * block_out); + rest = tail; + let start = b0; + s.spawn(move |_| { + Self::patchify_video_rgb_block_band( + frames, + width, + patch_size, + merge_size, + merged_patch, + pc_blocks, + start, + band, + lut, + ); + }); + b0 += nb; } - } + }); } + *out_idx = end_idx; Ok(()) } -} - -impl VisionPreProcessor for QwenVLProcessorBase { - fn default_mean(&self) -> [f64; 3] { - self.config.mean - } - - fn default_std(&self) -> [f64; 3] { - self.config.std - } - fn preprocess( + fn patchify_video_rgb_u8_chunk_into( &self, - images: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if images.is_empty() { - return Err(TransformError::EmptyBatch); + frames: &[VideoFrameRgb<'_>], + channel_order: RgbChannelOrder, + grid_h: usize, + grid_w: usize, + output: &mut [u8], + out_idx: &mut usize, + ) -> Result<(), TransformError> { + let patch_size = self.config.patch_size; + let merge_size = self.config.merge_size; + let temporal_patch_size = self.config.temporal_patch_size; + if frames.len() != temporal_patch_size { + return Err(TransformError::InvalidShape { + expected: format!("{temporal_patch_size} video frames in temporal patch"), + actual: vec![frames.len()], + }); } - // Store original sizes - let item_sizes: Vec<(u32, u32)> = images.iter().map(|img| img.dimensions()).collect(); + let height = grid_h * patch_size; + let width = grid_w * patch_size; + for frame in frames { + if frame.height != height || frame.width != width { + return Err(TransformError::InvalidShape { + expected: format!("video frame size {width}x{height}"), + actual: vec![frame.width, frame.height], + }); + } + } - let mean = config.get_image_mean(); - let std = config.get_image_std(); - // Qwen2VL/Qwen3VL image processors default to BICUBIC (PIL resample=3) - // when the preprocessor config omits `resample`. The global pil_to_filter - // fallback is bilinear, which yields smoother features and measurably - // degrades VLM accuracy, so pin the HF-correct default here. - let filter = pil_to_filter(config.resampling.or(Some(3))); + let merged_patch = merge_size * patch_size; + let pr_blocks = grid_h / merge_size; + let pc_blocks = grid_w / merge_size; + let n_blocks = pr_blocks * pc_blocks; + let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; + let base_idx = *out_idx; + let patch_values = n_blocks.checked_mul(block_out).ok_or_else(|| { + TransformError::ShapeError("Qwen deferred video patch size overflow".to_string()) + })?; + let end_idx = base_idx.checked_add(patch_values).ok_or_else(|| { + TransformError::ShapeError("Qwen deferred video patch range overflow".to_string()) + })?; + let region = output.get_mut(base_idx..end_idx).ok_or_else(|| { + TransformError::ShapeError("Qwen deferred video patch range out of bounds".to_string()) + })?; + let nthreads = par_threads(region.len(), n_blocks); + if nthreads <= 1 { + Self::patchify_video_rgb_u8_block_band( + frames, + channel_order, + width, + patch_size, + merge_size, + merged_patch, + pc_blocks, + 0, + region, + ); + } else { + let chunk_blocks = n_blocks.div_ceil(nthreads); + par_scope(|scope| { + let mut rest = &mut *region; + let mut block_start = 0; + while block_start < n_blocks { + let blocks = chunk_blocks.min(n_blocks - block_start); + let (band, tail) = rest.split_at_mut(blocks * block_out); + rest = tail; + let start = block_start; + scope.spawn(move |_| { + Self::patchify_video_rgb_u8_block_band( + frames, + channel_order, + width, + patch_size, + merge_size, + merged_patch, + pc_blocks, + start, + band, + ); + }); + block_start += blocks; + } + }); + } + *out_idx = end_idx; + Ok(()) + } + #[expect( + clippy::too_many_arguments, + reason = "resized RGB patchifier needs source order, grid, and output state" + )] + fn patchify_video_resized_u8_chunk_into( + &self, + frames: &[Cow<'_, [u8]>], + channel_order: RgbChannelOrder, + resize: &PilBicubicRgbPlan, + grid_h: usize, + grid_w: usize, + output: &mut [u8], + out_idx: &mut usize, + ) -> Result<(), TransformError> { let patch_size = self.config.patch_size; + let merge_size = self.config.merge_size; let temporal_patch_size = self.config.temporal_patch_size; - let patch_features = 3 * temporal_patch_size * patch_size * patch_size; - - // Pre-allocate based on total pixel count to avoid repeated Vec growth - let estimated_total: usize = images - .iter() - .map(|img| { - let (w, h) = img.dimensions(); - (w as usize * h as usize) / (self.config.merge_size * self.config.merge_size) - * patch_features - / (patch_size * patch_size) - }) - .sum(); - let mut all_patches: Vec = Vec::with_capacity(estimated_total); - let mut patches_per_image: Vec = Vec::with_capacity(images.len()); - let mut grid_thw_data = Vec::with_capacity(images.len() * 3); - let mut feature_token_counts = Vec::with_capacity(images.len()); - - for image in images { - let (w, h) = image.dimensions(); - let (target_h, target_w) = self.smart_resize(h as usize, w as usize)?; - - // Resize to the image's own target size (skip if dimensions match) - let (tw32, th32) = (target_w as u32, target_h as u32); - let needs_resize = config.do_resize.unwrap_or(true) && (w != tw32 || h != th32); - let resized; - let img_ref = if needs_resize { - // BICUBIC (Qwen default) must match PIL bit-for-bit so encoder - // inputs equal HF/vLLM; other filters keep the SIMD path. - resized = if filter == FilterType::CatmullRom { - resize_bicubic_pil(image, tw32, th32) - } else { - resize(image, tw32, th32, filter) - }; - &resized - } else { - image - }; - - // Grid dimensions based on the target size - let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_h, target_w, 1); - grid_thw_data.push(grid_t as i64); - grid_thw_data.push(grid_h as i64); - grid_thw_data.push(grid_w as i64); - - let num_patches = grid_t * grid_h * grid_w; - let tokens = self.calculate_tokens_from_grid(grid_t, grid_h, grid_w); - feature_token_counts.push(tokens); - - // Convert to tensor [C, H, W] and normalize in one fused pass - let tensor = if config.do_normalize.unwrap_or(true) { - to_tensor_and_normalize(img_ref, &mean, &std) - } else { - to_tensor(img_ref) - }; - - // Patchify directly into all_patches to avoid intermediate Vec + copy - self.patchify_into(&tensor, grid_t, grid_h, grid_w, &mut all_patches)?; - patches_per_image.push(num_patches as i64); + if frames.len() != temporal_patch_size { + return Err(TransformError::InvalidShape { + expected: format!("{temporal_patch_size} video frames in temporal patch"), + actual: vec![frames.len()], + }); } - let total_patches: usize = patches_per_image.iter().map(|&n| n as usize).sum(); - let encoder_input = - Array2::from_shape_vec((total_patches, patch_features), all_patches).map_err(|e| { - TransformError::ShapeError(format!( - "Failed to create patchified encoder_input [{total_patches}, {patch_features}]: {e}" - )) - })?; - - let result = PreprocessedEncoderInputs::new_dynamic( - encoder_input.into_dyn(), - feature_token_counts, - item_sizes, - ) - .with_extra( - "image_grid_thw", - ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), - ) - .with_extra( - "patches_per_image", - ModelSpecificValue::int_1d(patches_per_image), - ); - - Ok(result) + let merged_patch = merge_size * patch_size; + let row_blocks = grid_h / merge_size; + let column_blocks = grid_w / merge_size; + let blocks = row_blocks * column_blocks; + let values_per_block = + merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; + let values = blocks.checked_mul(values_per_block).ok_or_else(|| { + TransformError::ShapeError("Qwen resized video patch size overflow".to_string()) + })?; + let end = out_idx.checked_add(values).ok_or_else(|| { + TransformError::ShapeError("Qwen resized video patch range overflow".to_string()) + })?; + let region = output.get_mut(*out_idx..end).ok_or_else(|| { + TransformError::ShapeError("Qwen resized video patch range out of bounds".to_string()) + })?; + let workers = par_threads(region.len(), blocks); + let blocks_per_task = blocks.div_ceil(workers); + par_scope(|scope| { + for (task, band) in region + .chunks_mut(blocks_per_task * values_per_block) + .enumerate() + { + let block_start = task * blocks_per_task; + scope.spawn(move |_| { + Self::patchify_video_resized_u8_block_band( + frames, + channel_order, + resize, + patch_size, + merge_size, + merged_patch, + column_blocks, + block_start, + band, + ); + }); + } + }); + *out_idx = end; + Ok(()) } - fn preprocess_video( - &self, - frames: &[DynamicImage], - config: &PreProcessorConfig, - ) -> Result { - if frames.is_empty() { - return Err(TransformError::EmptyBatch); + #[expect( + clippy::too_many_arguments, + reason = "resized RGB patchifier mirrors the Qwen output layout" + )] + fn patchify_video_resized_u8_block_band( + frames: &[Cow<'_, [u8]>], + channel_order: RgbChannelOrder, + resize: &PilBicubicRgbPlan, + patch_size: usize, + merge_size: usize, + merged_patch: usize, + column_blocks: usize, + block_start: usize, + band: &mut [u8], + ) { + let values_per_block = merge_size * merge_size * 3 * frames.len() * patch_size * patch_size; + for (band_index, block_output) in band.chunks_mut(values_per_block).enumerate() { + let block = block_start + band_index; + let y0 = block / column_blocks * merged_patch; + let x0 = block % column_blocks * merged_patch; + let mut output_index = 0; + for merge_row in 0..merge_size { + for merge_column in 0..merge_size { + for channel in 0..3 { + let source_channel = channel_order.source_channel(channel); + for frame in frames { + for patch_row in 0..patch_size { + let y = y0 + merge_row * patch_size + patch_row; + let x = x0 + merge_column * patch_size; + for (patch_column, output) in block_output + [output_index..output_index + patch_size] + .iter_mut() + .enumerate() + { + *output = resize.pixel( + frame.as_ref(), + x + patch_column, + y, + source_channel, + ); + } + output_index += patch_size; + } + } + } + } + } } + } - let (w, h) = frames[0].dimensions(); - let item_sizes = vec![(w, h)]; - let mean = config.get_image_mean(); - let std = config.get_image_std(); - // Qwen2VL/Qwen3VL image processors default to BICUBIC (PIL resample=3) - // when the preprocessor config omits `resample`. The global pil_to_filter - // fallback is bilinear, which yields smoother features and measurably - // degrades VLM accuracy, so pin the HF-correct default here. - let filter = pil_to_filter(config.resampling.or(Some(3))); - - let temporal_patch_size = self.config.temporal_patch_size; - let padded_frames = frames.len().div_ceil(temporal_patch_size) * temporal_patch_size; - let (target_h, target_w) = self.smart_resize_video(frames.len(), h as usize, w as usize)?; - let (tw32, th32) = (target_w as u32, target_h as u32); + #[expect( + clippy::too_many_arguments, + reason = "deferred RGB patchifier mirrors the normalized output layout" + )] + fn patchify_video_rgb_u8_block_band( + frames: &[VideoFrameRgb<'_>], + channel_order: RgbChannelOrder, + width: usize, + patch_size: usize, + merge_size: usize, + merged_patch: usize, + pc_blocks: usize, + block_start: usize, + band: &mut [u8], + ) { + let block_out = merge_size * merge_size * 3 * frames.len() * patch_size * patch_size; + for (band_index, chunk) in band.chunks_mut(block_out).enumerate() { + let block = block_start + band_index; + let patch_row = block / pc_blocks; + let patch_column = block % pc_blocks; + let y0 = patch_row * merged_patch; + let x0 = patch_column * merged_patch; + let mut output_index = 0; + + for merge_row in 0..merge_size { + for merge_column in 0..merge_size { + for channel in 0..3 { + let source_channel = channel_order.source_channel(channel); + for frame in frames { + let raw = frame.data.as_ref(); + for patch_row in 0..patch_size { + let row = (y0 + merge_row * patch_size + patch_row) * width + + x0 + + merge_column * patch_size; + let mut source_index = row * 3 + source_channel; + for output in &mut chunk[output_index..output_index + patch_size] { + *output = raw[source_index]; + source_index += 3; + } + output_index += patch_size; + } + } + } + } + } + } + } + + #[expect( + clippy::too_many_arguments, + reason = "RGB video patchifier: frame window + grid dims + output band" + )] + fn patchify_video_rgb_block_band( + frames: &[VideoFrameRgb<'_>], + width: usize, + patch_size: usize, + merge_size: usize, + merged_patch: usize, + pc_blocks: usize, + block_start: usize, + band: &mut [f32], + lut: &[[f32; 256]; 3], + ) { + let block_out = merge_size * merge_size * 3 * frames.len() * patch_size * patch_size; + for (bi, chunk) in band.chunks_mut(block_out).enumerate() { + let blk = block_start + bi; + let pr = blk / pc_blocks; + let pc = blk % pc_blocks; + let y0 = pr * merged_patch; + let x0 = pc * merged_patch; + let mut o = 0usize; + + for mh in 0..merge_size { + for mw in 0..merge_size { + for (c, lut_c) in lut.iter().enumerate().take(3) { + for frame in frames { + let raw = frame.data.as_ref(); + for py in 0..patch_size { + let row = + (y0 + mh * patch_size + py) * width + x0 + mw * patch_size; + let mut src_idx = row * 3 + c; + for dst in &mut chunk[o..o + patch_size] { + *dst = lut_c[raw[src_idx] as usize]; + src_idx += 3; + } + o += patch_size; + } + } + } + } + } + } + } + + fn patchify_image_rgb_into( + &self, + image: &DynamicImage, + grid_h: usize, + grid_w: usize, + output: &mut [f32], + out_idx: &mut usize, + lut: &[[f32; 256]; 3], + ) -> Result<(), TransformError> { + let (width, height, data) = rgb_bytes(image); + let patch_size = self.config.patch_size; + let merge_size = self.config.merge_size; + let temporal_patch_size = self.config.temporal_patch_size; + let expected_height = grid_h * patch_size; + let expected_width = grid_w * patch_size; + if height != expected_height || width != expected_width { + return Err(TransformError::InvalidShape { + expected: format!("image size {expected_width}x{expected_height}"), + actual: vec![width, height], + }); + } + + let raw = data.as_ref(); + let merged_patch = merge_size * patch_size; + let pr_blocks = grid_h / merge_size; + let pc_blocks = grid_w / merge_size; + let n_blocks = pr_blocks * pc_blocks; + let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; + let base_idx = *out_idx; + let patch_values = n_blocks.checked_mul(block_out).ok_or_else(|| { + TransformError::ShapeError("Qwen image patch output size overflow".to_string()) + })?; + let end_idx = base_idx.checked_add(patch_values).ok_or_else(|| { + TransformError::ShapeError("Qwen image patch output range overflow".to_string()) + })?; + let region = output.get_mut(base_idx..end_idx).ok_or_else(|| { + TransformError::ShapeError("Qwen image patch output range out of bounds".to_string()) + })?; + let nthreads = par_threads(region.len() * 4, n_blocks); + if nthreads <= 1 { + Self::patchify_image_rgb_block_band( + raw, + width, + patch_size, + merge_size, + temporal_patch_size, + merged_patch, + pc_blocks, + 0, + region, + lut, + ); + } else { + let chunk_blocks = n_blocks.div_ceil(nthreads); + par_scope(|s| { + let mut rest = &mut *region; + let mut b0 = 0usize; + while b0 < n_blocks { + let nb = chunk_blocks.min(n_blocks - b0); + let (band, tail) = rest.split_at_mut(nb * block_out); + rest = tail; + let start = b0; + s.spawn(move |_| { + Self::patchify_image_rgb_block_band( + raw, + width, + patch_size, + merge_size, + temporal_patch_size, + merged_patch, + pc_blocks, + start, + band, + lut, + ); + }); + b0 += nb; + } + }); + } + *out_idx = end_idx; + + Ok(()) + } + + #[expect( + clippy::too_many_arguments, + reason = "RGB image patchifier: raw bytes + grid dims + output band" + )] + fn patchify_image_rgb_block_band( + raw: &[u8], + width: usize, + patch_size: usize, + merge_size: usize, + temporal_patch_size: usize, + merged_patch: usize, + pc_blocks: usize, + block_start: usize, + band: &mut [f32], + lut: &[[f32; 256]; 3], + ) { + let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; + for (bi, chunk) in band.chunks_mut(block_out).enumerate() { + let blk = block_start + bi; + let pr = blk / pc_blocks; + let pc = blk % pc_blocks; + let y0 = pr * merged_patch; + let x0 = pc * merged_patch; + let mut o = 0usize; + + for mh in 0..merge_size { + for mw in 0..merge_size { + for (c, lut_c) in lut.iter().enumerate().take(3) { + for _tp in 0..temporal_patch_size { + for py in 0..patch_size { + let row = + (y0 + mh * patch_size + py) * width + x0 + mw * patch_size; + let mut src_idx = row * 3 + c; + for dst in &mut chunk[o..o + patch_size] { + *dst = lut_c[raw[src_idx] as usize]; + src_idx += 3; + } + o += patch_size; + } + } + } + } + } + } + } + + fn preprocess_single_image_ref_deferred( + &self, + image: &DynamicImage, + config: &PreProcessorConfig, + ) -> Result { + let (width, height) = image.dimensions(); + let do_resize = config.do_resize.unwrap_or(true); + let (target_height, target_width) = if do_resize { + self.smart_resize(height as usize, width as usize)? + } else { + self.validate_unresized_patch_dimensions(height as usize, width as usize)?; + (height as usize, width as usize) + }; + let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_height, target_width, 1); + let patch_size = self.config.patch_size; + let temporal_patch_size = self.config.temporal_patch_size; + let patch_features = 3 * temporal_patch_size * patch_size * patch_size; + let num_patches = grid_t + .checked_mul(grid_h) + .and_then(|value| value.checked_mul(grid_w)) + .ok_or_else(|| { + TransformError::ShapeError("Qwen deferred image patch count overflow".to_string()) + })?; + let patch_values = num_patches.checked_mul(patch_features).ok_or_else(|| { + TransformError::ShapeError("Qwen deferred image patch size overflow".to_string()) + })?; + + let target_width = target_width as u32; + let target_height = target_height as u32; + let filter = pil_to_filter(config.resampling.or(Some(3))); + let resized; + let image = if do_resize && (width != target_width || height != target_height) { + resized = if filter == FilterType::CatmullRom { + resize_bicubic_pil(image, target_width, target_height) + } else { + resize(image, target_width, target_height, filter) + }; + &resized + } else { + image + }; + + let (rgb_width, rgb_height, rgb_data) = rgb_bytes(image); + let frames: Vec<_> = (0..temporal_patch_size) + .map(|_| VideoFrameRgb { + width: rgb_width, + height: rgb_height, + data: Cow::Borrowed(rgb_data.as_ref()), + }) + .collect(); + let mut patches = vec![0; patch_values]; + let mut output_index = 0; + self.patchify_video_rgb_u8_chunk_into( + &frames, + RgbChannelOrder::Rgb, + grid_h, + grid_w, + &mut patches, + &mut output_index, + )?; + debug_assert_eq!(output_index, patches.len()); + + let mean = config.get_image_mean(); + let std = config.get_image_std(); + let do_normalize = config.do_normalize.unwrap_or(true); + let scale: [f32; 3] = if do_normalize { + std::array::from_fn(|channel| 1.0 / (255.0 * std[channel] as f32)) + } else { + [1.0 / 255.0; 3] + }; + let bias: [f32; 3] = if do_normalize { + std::array::from_fn(|channel| -(mean[channel] as f32) / (std[channel] as f32)) + } else { + [0.0; 3] + }; + let lut: [[f32; 256]; 3] = std::array::from_fn(|channel| { + std::array::from_fn(|value| value as f32 * scale[channel] + bias[channel]) + }); + let channel_run = temporal_patch_size * patch_size * patch_size; + let deferred = DeferredNormalizedEncoderInput::new( + patches, + vec![num_patches, patch_features], + lut, + channel_run, + )?; + + Ok(PreprocessedEncoderInputs::new_deferred_normalized( + deferred, + vec![self.calculate_tokens_from_grid(grid_t, grid_h, grid_w)], + vec![(width, height)], + ) + .with_extra( + "image_grid_thw", + ModelSpecificValue::int_2d(vec![grid_t as i64, grid_h as i64, grid_w as i64], 1, 3), + ) + .with_extra( + "patches_per_image", + ModelSpecificValue::int_1d(vec![num_patches as i64]), + )) + } +} + +impl QwenVLProcessorBase { + fn preprocess_image_refs( + &self, + images: &[&DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + if images.is_empty() { + return Err(TransformError::EmptyBatch); + } + + let mean = config.get_image_mean(); + let std = config.get_image_std(); + // Qwen2VL/Qwen3VL image processors default to BICUBIC (PIL resample=3) + // when the preprocessor config omits `resample`. The global pil_to_filter + // fallback is bilinear, which yields smoother features and measurably + // degrades VLM accuracy, so pin the HF-correct default here. + let filter = pil_to_filter(config.resampling.or(Some(3))); + + let patch_size = self.config.patch_size; + let temporal_patch_size = self.config.temporal_patch_size; + let patch_features = 3 * temporal_patch_size * patch_size * patch_size; + let do_resize = config.do_resize.unwrap_or(true); + let do_normalize = config.do_normalize.unwrap_or(true); + let scale: [f32; 3] = if do_normalize { + std::array::from_fn(|c| 1.0 / (255.0 * std[c] as f32)) + } else { + [1.0 / 255.0; 3] + }; + let bias: [f32; 3] = if do_normalize { + std::array::from_fn(|c| -(mean[c] as f32) / (std[c] as f32)) + } else { + [0.0; 3] + }; + let lut: [[f32; 256]; 3] = + std::array::from_fn(|c| std::array::from_fn(|v| v as f32 * scale[c] + bias[c])); + + let mut image_plans = Vec::with_capacity(images.len()); + let mut item_sizes = Vec::with_capacity(images.len()); + let mut total_patch_values = 0usize; + let mut total_patches = 0usize; + for &image in images { + let (w, h) = image.dimensions(); + item_sizes.push((w, h)); + let (target_h, target_w) = if do_resize { + self.smart_resize(h as usize, w as usize)? + } else { + self.validate_unresized_patch_dimensions(h as usize, w as usize)?; + (h as usize, w as usize) + }; + let (tw32, th32) = (target_w as u32, target_h as u32); + let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_h, target_w, 1); + let num_patches = grid_t + .checked_mul(grid_h) + .and_then(|value| value.checked_mul(grid_w)) + .ok_or_else(|| { + TransformError::ShapeError(format!( + "Qwen image patch count overflow: grid=({grid_t}, {grid_h}, {grid_w})" + )) + })?; + total_patches = total_patches.checked_add(num_patches).ok_or_else(|| { + TransformError::ShapeError("Qwen image total patch count overflow".to_string()) + })?; + let patch_values = num_patches.checked_mul(patch_features).ok_or_else(|| { + TransformError::ShapeError(format!( + "Qwen image patch buffer size overflow: patches={num_patches}, features={patch_features}" + )) + })?; + total_patch_values = total_patch_values + .checked_add(patch_values) + .ok_or_else(|| { + TransformError::ShapeError( + "Qwen image patch buffer total size overflow".to_string(), + ) + })?; + image_plans.push(QwenImagePlan { + target_width: tw32, + target_height: th32, + needs_resize: do_resize && (w != tw32 || h != th32), + grid_t, + grid_h, + grid_w, + num_patches, + patch_values, + tokens: self.calculate_tokens_from_grid(grid_t, grid_h, grid_w), + }); + } + + let mut all_patches: Vec = Vec::with_capacity(total_patch_values); + let mut patches_per_image: Vec = Vec::with_capacity(images.len()); + let mut grid_thw_data = Vec::with_capacity(images.len() * 3); + let mut feature_token_counts = Vec::with_capacity(images.len()); + + for (image, plan) in images.iter().copied().zip(image_plans) { + // Resize to the image's own target size (skip if dimensions match) + let resized; + let img_ref = if plan.needs_resize { + // BICUBIC (Qwen default) uses the PIL-compatible path; other + // filters keep the SIMD path. + resized = if filter == FilterType::CatmullRom { + resize_bicubic_pil(image, plan.target_width, plan.target_height) + } else { + resize(image, plan.target_width, plan.target_height, filter) + }; + &resized + } else { + image + }; + + grid_thw_data.push(plan.grid_t as i64); + grid_thw_data.push(plan.grid_h as i64); + grid_thw_data.push(plan.grid_w as i64); + + feature_token_counts.push(plan.tokens); + + // Patchify directly from RGB bytes to avoid the intermediate + // [C,H,W] tensor allocation. This matches the tensor path's + // channel/temporal/spatial order. + let base_idx = all_patches.len(); + all_patches.resize(base_idx + plan.patch_values, 0.0); + let mut out_idx = base_idx; + self.patchify_image_rgb_into( + img_ref, + plan.grid_h, + plan.grid_w, + &mut all_patches, + &mut out_idx, + &lut, + )?; + debug_assert_eq!(out_idx, all_patches.len()); + patches_per_image.push(plan.num_patches as i64); + } + + let encoder_input = + Array2::from_shape_vec((total_patches, patch_features), all_patches).map_err(|e| { + TransformError::ShapeError(format!( + "Failed to create patchified encoder_input [{total_patches}, {patch_features}]: {e}" + )) + })?; + + let result = PreprocessedEncoderInputs::new_dynamic( + encoder_input.into_dyn(), + feature_token_counts, + item_sizes, + ) + .with_extra( + "image_grid_thw", + ModelSpecificValue::int_2d(grid_thw_data, images.len(), 3), + ) + .with_extra( + "patches_per_image", + ModelSpecificValue::int_1d(patches_per_image), + ); + + Ok(result) + } + + fn preprocess_image_refs_deferred( + &self, + images: &[&DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + if let [image] = images { + self.preprocess_single_image_ref_deferred(image, config) + } else { + self.preprocess_image_refs(images, config) + } + } + + fn preprocess_video( + &self, + frames: &[DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + if frames.is_empty() { + return Err(TransformError::EmptyBatch); + } + + let (w, h) = frames[0].dimensions(); + let item_sizes = vec![(w, h)]; + let mean = config.get_image_mean(); + let std = config.get_image_std(); + // Qwen2VL/Qwen3VL image processors default to BICUBIC (PIL resample=3) + // when the preprocessor config omits `resample`. The global pil_to_filter + // fallback is bilinear, which yields smoother features and measurably + // degrades VLM accuracy, so pin the HF-correct default here. + let filter = pil_to_filter(config.resampling.or(Some(3))); + + let temporal_patch_size = self.config.temporal_patch_size; + let padded_frames = frames.len().div_ceil(temporal_patch_size) * temporal_patch_size; + let do_resize = config.do_resize.unwrap_or(true); + let (target_h, target_w) = if do_resize { + self.smart_resize_video(frames.len(), h as usize, w as usize)? + } else { + self.validate_unresized_patch_dimensions(h as usize, w as usize)?; + (h as usize, w as usize) + }; + let (tw32, th32) = (target_w as u32, target_h as u32); let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_h, target_w, padded_frames); let patch_size = self.config.patch_size; @@ -759,22 +1495,14 @@ impl VisionPreProcessor for QwenVLProcessorBase { for tp in 0..temporal_patch_size { let idx = (gt * temporal_patch_size + tp).min(frames.len() - 1); let frame = &frames[idx]; - let needs_resize = config.do_resize.unwrap_or(true) - && (frame.width() != tw32 || frame.height() != th32); + let needs_resize = do_resize && (frame.width() != tw32 || frame.height() != th32); if needs_resize { - // BICUBIC (Qwen default) must match PIL bit-for-bit so video - // encoder inputs equal HF/vLLM, same as the image path; other - // filters keep the SIMD resizer. - let resized = if filter == FilterType::CatmullRom { - resize_bicubic_pil(frame, tw32, th32) - } else { - resize(frame, tw32, th32, filter) - }; - let (width, height, data) = rgb_bytes(&resized); + let (width, height, data) = + resize_dynamic_frame_to_raw(frame, tw32, th32, filter); frame_rgbs.push(VideoFrameRgb { width, height, - data: Cow::Owned(data.into_owned()), + data: Cow::Owned(data), }); } else { let (width, height, data) = rgb_bytes(frame); @@ -847,7 +1575,13 @@ impl VisionPreProcessor for QwenVLProcessorBase { let temporal_patch_size = self.config.temporal_patch_size; let padded_frames = frames.len().div_ceil(temporal_patch_size) * temporal_patch_size; - let (target_h, target_w) = self.smart_resize_video(frames.len(), h as usize, w as usize)?; + let do_resize = config.do_resize.unwrap_or(true); + let (target_h, target_w) = if do_resize { + self.smart_resize_video(frames.len(), h as usize, w as usize)? + } else { + self.validate_unresized_patch_dimensions(h as usize, w as usize)?; + (h as usize, w as usize) + }; let (tw32, th32) = (target_w as u32, target_h as u32); let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_h, target_w, padded_frames); @@ -871,8 +1605,6 @@ impl VisionPreProcessor for QwenVLProcessorBase { let lut: [[f32; 256]; 3] = std::array::from_fn(|c| std::array::from_fn(|v| v as f32 * scale[c] + bias[c])); - let mut needs_resize_any = false; - let do_resize = config.do_resize.unwrap_or(true); for frame in frames { let expected_len = (frame.width as usize) .checked_mul(frame.height as usize) @@ -892,86 +1624,38 @@ impl VisionPreProcessor for QwenVLProcessorBase { actual: vec![frame.data.len()], }); } - needs_resize_any |= do_resize && (frame.width != tw32 || frame.height != th32); } let mut out_idx = 0; - if needs_resize_any { - let mut frame_rgbs = Vec::with_capacity(temporal_patch_size); - for gt in 0..grid_t { - frame_rgbs.clear(); - for tp in 0..temporal_patch_size { - let idx = (gt * temporal_patch_size + tp).min(frames.len() - 1); - let frame = frames[idx]; - let needs_resize = do_resize && (frame.width != tw32 || frame.height != th32); - if needs_resize { - // BICUBIC (Qwen default) must match PIL bit-for-bit so video - // encoder inputs equal HF/vLLM, same as the image path; other - // filters keep the SIMD resizer. - let resized = if filter == FilterType::CatmullRom { - resize_bicubic_pil_rgb( - frame.data, - frame.width, - frame.height, - tw32, - th32, - )? - } else { - resize_rgb_bytes( - frame.data, - frame.width, - frame.height, - tw32, - th32, - filter, - )? - }; - frame_rgbs.push(VideoFrameRgb { - width: tw32 as usize, - height: th32 as usize, - data: Cow::Owned(resized.into_raw()), - }); - } else { - frame_rgbs.push(VideoFrameRgb { - width: frame.width as usize, - height: frame.height as usize, - data: Cow::Borrowed(frame.data), - }); - } - } - - self.patchify_video_rgb_chunk_into( - &frame_rgbs, - grid_h, - grid_w, - &mut all_patches, - &mut out_idx, - &lut, - )?; - } - } else { - let mut frame_rgbs = Vec::with_capacity(temporal_patch_size); - for gt in 0..grid_t { - frame_rgbs.clear(); - for tp in 0..temporal_patch_size { - let idx = (gt * temporal_patch_size + tp).min(frames.len() - 1); - let frame = frames[idx]; + let mut frame_rgbs = Vec::with_capacity(temporal_patch_size); + for gt in 0..grid_t { + frame_rgbs.clear(); + for tp in 0..temporal_patch_size { + let idx = (gt * temporal_patch_size + tp).min(frames.len() - 1); + let frame = frames[idx]; + if do_resize && (frame.width != tw32 || frame.height != th32) { + frame_rgbs.push(VideoFrameRgb { + width: tw32 as usize, + height: th32 as usize, + data: Cow::Owned(resize_rgb_frame_to_raw(frame, tw32, th32, filter)?), + }); + } else { frame_rgbs.push(VideoFrameRgb { width: frame.width as usize, height: frame.height as usize, data: Cow::Borrowed(frame.data), }); } - - self.patchify_video_rgb_chunk_into( - &frame_rgbs, - grid_h, - grid_w, - &mut all_patches, - &mut out_idx, - &lut, - )?; } + + self.patchify_video_rgb_chunk_into( + &frame_rgbs, + grid_h, + grid_w, + &mut all_patches, + &mut out_idx, + &lut, + )?; } debug_assert_eq!(out_idx, all_patches.len()); @@ -1000,17 +1684,424 @@ impl VisionPreProcessor for QwenVLProcessorBase { ModelSpecificValue::int_1d(vec![num_patches as i64]), ); - Ok(result) + Ok(result) + } + + fn preprocess_video_rgb_deferred( + &self, + frames: &[RgbFrameRef<'_>], + config: &PreProcessorConfig, + ) -> Result { + if frames.is_empty() { + return Err(TransformError::EmptyBatch); + } + + let width = frames[0].width; + let height = frames[0].height; + let item_sizes = vec![(width, height)]; + let mean = config.get_image_mean(); + let std = config.get_image_std(); + let filter = pil_to_filter(config.resampling.or(Some(3))); + let temporal_patch_size = self.config.temporal_patch_size; + let padded_frames = frames.len().div_ceil(temporal_patch_size) * temporal_patch_size; + let do_resize = config.do_resize.unwrap_or(true); + let (target_h, target_w) = if do_resize { + self.smart_resize_video(frames.len(), height as usize, width as usize)? + } else { + self.validate_unresized_patch_dimensions(height as usize, width as usize)?; + (height as usize, width as usize) + }; + let (target_width, target_height) = (target_w as u32, target_h as u32); + let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_h, target_w, padded_frames); + let patch_size = self.config.patch_size; + let patch_features = 3 * temporal_patch_size * patch_size * patch_size; + let num_patches = grid_t * grid_h * grid_w; + let tokens = self.calculate_tokens_from_grid(grid_t, grid_h, grid_w); + let mut patches = vec![0; num_patches * patch_features]; + + let do_normalize = config.do_normalize.unwrap_or(true); + let scale: [f32; 3] = if do_normalize { + std::array::from_fn(|channel| 1.0 / (255.0 * std[channel] as f32)) + } else { + [1.0 / 255.0; 3] + }; + let bias: [f32; 3] = if do_normalize { + std::array::from_fn(|channel| -(mean[channel] as f32) / (std[channel] as f32)) + } else { + [0.0; 3] + }; + let lut: [[f32; 256]; 3] = std::array::from_fn(|channel| { + std::array::from_fn(|value| value as f32 * scale[channel] + bias[channel]) + }); + + for frame in frames { + let expected_len = (frame.width as usize) + .checked_mul(frame.height as usize) + .and_then(|pixels| pixels.checked_mul(3)) + .ok_or_else(|| { + TransformError::ShapeError(format!( + "video frame dimensions are too large: {}x{}", + frame.width, frame.height + )) + })?; + if frame.data.len() != expected_len { + return Err(TransformError::InvalidShape { + expected: format!( + "RGB frame byte length {expected_len} for {}x{}", + frame.width, frame.height + ), + actual: vec![frame.data.len()], + }); + } + } + + let direct_resize = (do_resize + && filter == FilterType::CatmullRom + && (width != target_width || height != target_height) + && frames + .iter() + .all(|frame| frame.width == width && frame.height == height)) + .then(|| PilBicubicRgbPlan::new(width, height, target_width, target_height)) + .transpose()?; + + let mut output_index = 0; + let mut frame_rgbs = Vec::with_capacity(temporal_patch_size); + let mut resized_horizontal = Vec::with_capacity(temporal_patch_size); + for temporal_index in 0..grid_t { + if let Some(resize) = &direct_resize { + resized_horizontal.clear(); + for temporal_patch_index in 0..temporal_patch_size { + let frame_index = (temporal_index * temporal_patch_size + temporal_patch_index) + .min(frames.len() - 1); + resized_horizontal.push(resize.prepare_horizontal(frames[frame_index].data)?); + } + self.patchify_video_resized_u8_chunk_into( + &resized_horizontal, + RgbChannelOrder::Rgb, + resize, + grid_h, + grid_w, + &mut patches, + &mut output_index, + )?; + continue; + } + + frame_rgbs.clear(); + for temporal_patch_index in 0..temporal_patch_size { + let frame_index = (temporal_index * temporal_patch_size + temporal_patch_index) + .min(frames.len() - 1); + let frame = frames[frame_index]; + if do_resize && (frame.width != target_width || frame.height != target_height) { + frame_rgbs.push(VideoFrameRgb { + width: target_width as usize, + height: target_height as usize, + data: Cow::Owned(resize_rgb_frame_to_raw( + frame, + target_width, + target_height, + filter, + )?), + }); + } else { + frame_rgbs.push(VideoFrameRgb { + width: frame.width as usize, + height: frame.height as usize, + data: Cow::Borrowed(frame.data), + }); + } + } + self.patchify_video_rgb_u8_chunk_into( + &frame_rgbs, + RgbChannelOrder::Rgb, + grid_h, + grid_w, + &mut patches, + &mut output_index, + )?; + } + debug_assert_eq!(output_index, patches.len()); + + let channel_run = temporal_patch_size * patch_size * patch_size; + let deferred = DeferredNormalizedEncoderInput::new( + patches, + vec![num_patches, patch_features], + lut, + channel_run, + )?; + Ok( + PreprocessedEncoderInputs::new_deferred_normalized(deferred, vec![tokens], item_sizes) + .with_extra( + "video_grid_thw", + ModelSpecificValue::int_2d( + vec![grid_t as i64, grid_h as i64, grid_w as i64], + 1, + 3, + ), + ) + .with_extra( + "patches_per_video", + ModelSpecificValue::int_1d(vec![num_patches as i64]), + ) + .with_extra( + "patches_per_image", + ModelSpecificValue::int_1d(vec![num_patches as i64]), + ), + ) + } + + fn preprocess_video_rgb_stream_deferred( + &self, + stream: DecodedRgbFrameStream, + config: &PreProcessorConfig, + ) -> Result { + let frame_count = stream.expected_frames(); + if frame_count == 0 { + return Err(TransformError::EmptyBatch); + } + let first = stream + .next_frame() + .map_err(TransformError::ShapeError)? + .ok_or_else(|| { + TransformError::ShapeError( + "decoded RGB stream ended before its first frame".to_string(), + ) + })?; + let width = first.width; + let height = first.height; + let channel_order = first.channel_order; + let item_sizes = vec![(width, height)]; + let temporal_patch_size = self.config.temporal_patch_size; + let padded_frames = frame_count.div_ceil(temporal_patch_size) * temporal_patch_size; + let do_resize = config.do_resize.unwrap_or(true); + let (target_h, target_w) = if do_resize { + self.smart_resize_video(frame_count, height as usize, width as usize)? + } else { + self.validate_unresized_patch_dimensions(height as usize, width as usize)?; + (height as usize, width as usize) + }; + let (target_width, target_height) = (target_w as u32, target_h as u32); + let (grid_t, grid_h, grid_w) = self.calculate_grid_thw(target_h, target_w, padded_frames); + let patch_size = self.config.patch_size; + let patch_features = 3 * temporal_patch_size * patch_size * patch_size; + let num_patches = grid_t * grid_h * grid_w; + let tokens = self.calculate_tokens_from_grid(grid_t, grid_h, grid_w); + let mut patches = vec![0; num_patches * patch_features]; + + let mean = config.get_image_mean(); + let std = config.get_image_std(); + let do_normalize = config.do_normalize.unwrap_or(true); + let scale: [f32; 3] = if do_normalize { + std::array::from_fn(|channel| 1.0 / (255.0 * std[channel] as f32)) + } else { + [1.0 / 255.0; 3] + }; + let bias: [f32; 3] = if do_normalize { + std::array::from_fn(|channel| -(mean[channel] as f32) / (std[channel] as f32)) + } else { + [0.0; 3] + }; + let lut: [[f32; 256]; 3] = std::array::from_fn(|channel| { + std::array::from_fn(|value| value as f32 * scale[channel] + bias[channel]) + }); + let filter = pil_to_filter(config.resampling.or(Some(3))); + let direct_resize = (do_resize + && filter == FilterType::CatmullRom + && (width != target_width || height != target_height)) + .then(|| PilBicubicRgbPlan::new(width, height, target_width, target_height)) + .transpose()?; + + let mut consumed = 1usize; + let mut output_frame_index = 0usize; + let mut last_frame = first; + let mut output_index = 0; + let mut frame_group = Vec::with_capacity(temporal_patch_size); + for _ in 0..grid_t { + frame_group.clear(); + for _ in 0..temporal_patch_size { + if output_frame_index == 0 { + frame_group.push(last_frame.clone()); + output_frame_index += 1; + continue; + } + if output_frame_index < frame_count { + last_frame = stream + .next_frame() + .map_err(TransformError::ShapeError)? + .ok_or_else(|| { + TransformError::ShapeError(format!( + "decoded RGB stream ended after {consumed} of {frame_count} frames" + )) + })?; + consumed += 1; + } + frame_group.push(last_frame.clone()); + output_frame_index += 1; + } + + for frame in &frame_group { + validate_streamed_qwen_frame(frame, width, height, channel_order)?; + } + if let Some(resize) = &direct_resize { + let mut resized_horizontal = Vec::with_capacity(temporal_patch_size); + for frame in &frame_group { + resized_horizontal.push(resize.prepare_horizontal(frame.data.as_ref())?); + } + self.patchify_video_resized_u8_chunk_into( + &resized_horizontal, + channel_order, + resize, + grid_h, + grid_w, + &mut patches, + &mut output_index, + )?; + continue; + } + + let mut frame_rgbs = Vec::with_capacity(temporal_patch_size); + for frame in &frame_group { + let frame_ref = RgbFrameRef { + width: frame.width, + height: frame.height, + data: frame.data.as_ref(), + }; + if do_resize && (frame.width != target_width || frame.height != target_height) { + frame_rgbs.push(VideoFrameRgb { + width: target_width as usize, + height: target_height as usize, + data: Cow::Owned(resize_rgb_frame_to_raw( + frame_ref, + target_width, + target_height, + filter, + )?), + }); + } else { + frame_rgbs.push(VideoFrameRgb { + width: frame.width as usize, + height: frame.height as usize, + data: Cow::Borrowed(frame.data.as_ref()), + }); + } + } + self.patchify_video_rgb_u8_chunk_into( + &frame_rgbs, + channel_order, + grid_h, + grid_w, + &mut patches, + &mut output_index, + )?; + } + if consumed != frame_count { + return Err(TransformError::ShapeError(format!( + "decoded RGB stream consumed {consumed} frames, expected {frame_count}" + ))); + } + if stream + .next_frame() + .map_err(TransformError::ShapeError)? + .is_some() + { + return Err(TransformError::ShapeError( + "decoded RGB stream produced more frames than expected".to_string(), + )); + } + debug_assert_eq!(output_index, patches.len()); + + let channel_run = temporal_patch_size * patch_size * patch_size; + let deferred = DeferredNormalizedEncoderInput::new( + patches, + vec![num_patches, patch_features], + lut, + channel_run, + )?; + Ok( + PreprocessedEncoderInputs::new_deferred_normalized(deferred, vec![tokens], item_sizes) + .with_extra( + "video_grid_thw", + ModelSpecificValue::int_2d( + vec![grid_t as i64, grid_h as i64, grid_w as i64], + 1, + 3, + ), + ) + .with_extra( + "patches_per_video", + ModelSpecificValue::int_1d(vec![num_patches as i64]), + ) + .with_extra( + "patches_per_image", + ModelSpecificValue::int_1d(vec![num_patches as i64]), + ), + ) + } +} + +impl VisionPreProcessor for QwenVLProcessorBase { + fn default_mean(&self) -> [f64; 3] { + self.config.mean + } + + fn default_std(&self) -> [f64; 3] { + self.config.std + } + + fn preprocess( + &self, + images: &[DynamicImage], + config: &PreProcessorConfig, + ) -> Result { + if images.len() == 1 { + let image_refs = [&images[0]]; + return self.preprocess_image_refs(&image_refs, config); + } + let image_refs = images.iter().collect::>(); + self.preprocess_image_refs(&image_refs, config) + } + + fn preprocess_vision_input( + &self, + request: VisionPreprocessRequest<'_>, + config: &PreProcessorConfig, + ) -> Result { + match request.input { + VisionInput::Images(images) => match request.output { + OutputPreference::Materialized => self.preprocess_image_refs(images, config), + OutputPreference::CompactAllowed => { + self.preprocess_image_refs_deferred(images, config) + } + }, + VisionInput::Video(VideoInput::Frames(frames)) => self.preprocess_video(frames, config), + VisionInput::Video(VideoInput::Rgb(frames)) => match request.output { + OutputPreference::Materialized => self.preprocess_video_rgb(frames, config), + OutputPreference::CompactAllowed => { + self.preprocess_video_rgb_deferred(frames, config) + } + }, + VisionInput::Video(VideoInput::RgbStream(stream)) => { + let mut output = self.preprocess_video_rgb_stream_deferred(stream, config)?; + if request.output == OutputPreference::Materialized { + output.materialize_encoder_input()?; + } + Ok(output) + } + } } - fn calculate_num_tokens(&self, width: u32, height: u32, _config: &PreProcessorConfig) -> usize { - // Calculate resized dimensions - let (new_height, new_width) = match self.smart_resize(height as usize, width as usize) { - Ok((h, w)) => (h, w), - Err(_) => { - // Fallback: use minimum size - let factor = self.get_factor(); - (factor, factor) + fn calculate_num_tokens(&self, width: u32, height: u32, config: &PreProcessorConfig) -> usize { + let fallback_size = || { + let factor = self.get_factor(); + (factor, factor) + }; + let (new_height, new_width) = if config.do_resize.unwrap_or(true) { + self.smart_resize(height as usize, width as usize) + .unwrap_or_else(|_| fallback_size()) + } else { + match self.validate_unresized_patch_dimensions(height as usize, width as usize) { + Ok(()) => (height as usize, width as usize), + Err(_) => fallback_size(), } }; @@ -1031,9 +2122,15 @@ impl VisionPreProcessor for QwenVLProcessorBase { #[cfg(test)] mod tests { + use std::{sync::mpsc::sync_channel, thread}; + use image::RgbImage; use super::*; + use crate::vision::{ + processor::{ModalityPreProcessor, PreprocessRequest}, + transforms::to_tensor_and_normalize, + }; fn create_test_config() -> QwenVLConfig { QwenVLConfig { @@ -1097,12 +2194,152 @@ mod tests { DynamicImage::ImageRgb8(image) } + #[test] + fn test_calculate_num_tokens_honors_do_resize_false() { + let processor = QwenVLProcessorBase::new(create_test_config()); + let config = PreProcessorConfig { + do_resize: Some(false), + ..Default::default() + }; + + let tokens = processor.calculate_num_tokens(84, 56, &config); + let (grid_t, grid_h, grid_w) = processor.calculate_grid_thw(56, 84, 1); + + assert_eq!( + tokens, + processor.calculate_tokens_from_grid(grid_t, grid_h, grid_w) + ); + assert_ne!( + tokens, + processor.calculate_num_tokens(84, 56, &PreProcessorConfig::default()) + ); + } + + #[test] + fn test_preprocess_image_matches_tensor_patchify_with_resize() { + let processor = QwenVLProcessorBase::new(create_video_test_config()); + let config = PreProcessorConfig { + image_mean: Some(processor.default_mean().to_vec()), + image_std: Some(processor.default_std().to_vec()), + ..Default::default() + }; + let image = create_sized_pattern_frame(7, 9, 3); + let (target_h, target_w) = processor.smart_resize(9, 7).unwrap(); + assert!( + (target_w as u32, target_h as u32) != (7u32, 9u32), + "test must force a resize; target {target_w}x{target_h} should differ from 7x9" + ); + + let result = processor + .preprocess(std::slice::from_ref(&image), &config) + .unwrap(); + let actual = result.encoder_input.as_slice_memory_order().unwrap(); + + let resized = resize_bicubic_pil(&image, target_w as u32, target_h as u32); + let tensor = to_tensor_and_normalize( + &resized, + &processor.default_mean(), + &processor.default_std(), + ); + let (grid_t, grid_h, grid_w) = processor.calculate_grid_thw(target_h, target_w, 1); + let mut expected = Vec::new(); + processor + .patchify_into(&tensor, grid_t, grid_h, grid_w, &mut expected) + .unwrap(); + + assert_eq!(actual.len(), expected.len()); + for (idx, (&got, &want)) in actual.iter().zip(expected.iter()).enumerate() { + assert_eq!( + got.to_bits(), + want.to_bits(), + "image patch value differs at index {idx}: got {got}, want {want}" + ); + } + } + + #[test] + fn test_patchify_image_rgb_block_band_matches_tensor_patchify_with_resize() { + let processor = QwenVLProcessorBase::new(create_video_test_config()); + let image = create_sized_pattern_frame(7, 9, 3); + let (target_h, target_w) = processor.smart_resize(9, 7).unwrap(); + let resized = resize_bicubic_pil(&image, target_w as u32, target_h as u32); + let tensor = to_tensor_and_normalize( + &resized, + &processor.default_mean(), + &processor.default_std(), + ); + let (grid_t, grid_h, grid_w) = processor.calculate_grid_thw(target_h, target_w, 1); + let mut expected = Vec::new(); + processor + .patchify_into(&tensor, grid_t, grid_h, grid_w, &mut expected) + .unwrap(); + + let (width, height, raw) = rgb_bytes(&resized); + assert_eq!((width, height), (target_w, target_h)); + let patch_size = processor.config.patch_size; + let merge_size = processor.config.merge_size; + let temporal_patch_size = processor.config.temporal_patch_size; + let merged_patch = merge_size * patch_size; + let pr_blocks = grid_h / merge_size; + let pc_blocks = grid_w / merge_size; + let n_blocks = pr_blocks * pc_blocks; + assert!( + n_blocks > 1, + "test must exercise multiple image patch blocks" + ); + let block_out = merge_size * merge_size * 3 * temporal_patch_size * patch_size * patch_size; + let mean = processor.default_mean(); + let std = processor.default_std(); + let scale: [f32; 3] = std::array::from_fn(|c| 1.0 / (255.0 * std[c] as f32)); + let bias: [f32; 3] = std::array::from_fn(|c| -(mean[c] as f32) / (std[c] as f32)); + let lut: [[f32; 256]; 3] = + std::array::from_fn(|c| std::array::from_fn(|v| v as f32 * scale[c] + bias[c])); + + let mut actual = vec![0.0; expected.len()]; + let split_blocks = n_blocks / 2; + let split_at = split_blocks * block_out; + let (first, second) = actual.split_at_mut(split_at); + QwenVLProcessorBase::patchify_image_rgb_block_band( + raw.as_ref(), + width, + patch_size, + merge_size, + temporal_patch_size, + merged_patch, + pc_blocks, + 0, + first, + &lut, + ); + QwenVLProcessorBase::patchify_image_rgb_block_band( + raw.as_ref(), + width, + patch_size, + merge_size, + temporal_patch_size, + merged_patch, + pc_blocks, + split_blocks, + second, + &lut, + ); + + assert_eq!(actual.len(), expected.len()); + for (idx, (&got, &want)) in actual.iter().zip(expected.iter()).enumerate() { + assert_eq!( + got.to_bits(), + want.to_bits(), + "image block-band patch value differs at index {idx}: got {got}, want {want}" + ); + } + } + /// When a video frame actually needs resizing, the DynamicImage path /// (`preprocess_video` → `resize_bicubic_pil`) and the raw-RGB path /// (`preprocess_video_rgb` → `resize_bicubic_pil_rgb`) must produce /// byte-for-byte identical encoder inputs. The other video tests use 4x4 - /// frames that need no resize, so this is the only one exercising the - /// default-bicubic resize branch added for HF/vLLM parity. + /// frames that need no resize, so this is the one exercising the + /// default-bicubic resize branch. #[test] fn test_preprocess_video_rgb_matches_dynamic_with_resize() { let processor = QwenVLProcessorBase::new(create_video_test_config()); @@ -1158,6 +2395,233 @@ mod tests { } } + #[test] + fn test_preprocess_image_deferred_matches_fp32_and_bf16_reference() { + let processor = QwenVLProcessorBase::new(create_video_test_config()); + let config = PreProcessorConfig { + image_mean: Some(processor.default_mean().to_vec()), + image_std: Some(processor.default_std().to_vec()), + ..Default::default() + }; + let image = create_sized_pattern_frame(7, 9, 31); + let reference = processor.preprocess_image_refs(&[&image], &config).unwrap(); + let mut deferred = processor + .preprocess_image_refs_deferred(&[&image], &config) + .unwrap(); + + let reference_values = reference.encoder_input.as_slice_memory_order().unwrap(); + let compact = deferred.encoder_input.deferred_normalized().unwrap(); + let mut actual_bf16 = vec![0; compact.len() * 2]; + compact.fill_bf16_le_bytes(&mut actual_bf16).unwrap(); + let reference_bf16 = reference_values + .iter() + .flat_map(|&value| { + let bits = value.to_bits(); + let lsb = (bits >> 16) & 1; + ((bits.wrapping_add(0x7fff + lsb) >> 16) as u16).to_le_bytes() + }) + .collect::>(); + assert_eq!(actual_bf16, reference_bf16); + + deferred.materialize_encoder_input().unwrap(); + let actual_values = deferred.encoder_input.as_slice_memory_order().unwrap(); + assert_eq!(actual_values.len(), reference_values.len()); + for (index, (&actual, &expected)) in actual_values.iter().zip(reference_values).enumerate() + { + assert_eq!( + actual.to_bits(), + expected.to_bits(), + "deferred image path diverges at index {index}" + ); + } + } + + #[test] + fn test_preprocess_video_rgb_matches_dynamic_with_resize_and_padding() { + let processor = QwenVLProcessorBase::new(create_video_test_config()); + let config = PreProcessorConfig { + image_mean: Some(processor.default_mean().to_vec()), + image_std: Some(processor.default_std().to_vec()), + ..Default::default() + }; + let frames = vec![ + create_sized_pattern_frame(7, 9, 3), + create_sized_pattern_frame(7, 9, 101), + create_sized_pattern_frame(7, 9, 177), + ]; + assert_ne!( + frames.len() % processor.temporal_patch_size(), + 0, + "test must force temporal padding" + ); + let (target_h, target_w) = processor.smart_resize_video(frames.len(), 9, 7).unwrap(); + assert!( + (target_w as u32, target_h as u32) != (7u32, 9u32), + "test must force a resize; target {target_w}x{target_h} should differ from 7x9" + ); + + let rgb_frames = frames + .iter() + .map(|frame| { + let DynamicImage::ImageRgb8(rgb) = frame else { + panic!("test frame is not RGB8"); + }; + RgbFrameRef { + width: rgb.width(), + height: rgb.height(), + data: rgb.as_raw(), + } + }) + .collect::>(); + + let dynamic = processor.preprocess_video(&frames, &config).unwrap(); + let rgb = processor + .preprocess_video_rgb(&rgb_frames, &config) + .unwrap(); + let mut deferred = processor + .preprocess_video_rgb_deferred(&rgb_frames, &config) + .unwrap(); + + let a = dynamic.encoder_input.as_slice_memory_order().unwrap(); + let b = rgb.encoder_input.as_slice_memory_order().unwrap(); + let compact = deferred.encoder_input.deferred_normalized().unwrap(); + let mut deferred_bf16 = vec![0; compact.len() * 2]; + compact.fill_bf16_le_bytes(&mut deferred_bf16).unwrap(); + let reference_bf16 = b + .iter() + .flat_map(|&value| { + let bits = value.to_bits(); + let lsb = (bits >> 16) & 1; + ((bits.wrapping_add(0x7fff + lsb) >> 16) as u16).to_le_bytes() + }) + .collect::>(); + assert_eq!(deferred_bf16, reference_bf16); + deferred.materialize_encoder_input().unwrap(); + let c = deferred.encoder_input.as_slice_memory_order().unwrap(); + assert_eq!(a.len(), b.len()); + assert_eq!(b.len(), c.len()); + for (idx, ((&got, &want), &deferred_value)) in + a.iter().zip(b.iter()).zip(c.iter()).enumerate() + { + assert_eq!( + got.to_bits(), + want.to_bits(), + "resized padded video path diverges at index {idx}: dynamic {got} vs rgb {want}" + ); + assert_eq!( + want.to_bits(), + deferred_value.to_bits(), + "deferred video path diverges at index {idx}: rgb {want} vs deferred {deferred_value}" + ); + } + } + + fn assert_bgr_streaming_matches_bulk(resampling: Option) { + let processor = QwenVLProcessorBase::new(create_video_test_config()); + let config = PreProcessorConfig { + image_mean: Some(processor.default_mean().to_vec()), + image_std: Some(processor.default_std().to_vec()), + resampling, + ..Default::default() + }; + let frames = [ + create_sized_pattern_frame(7, 9, 3), + create_sized_pattern_frame(7, 9, 101), + create_sized_pattern_frame(7, 9, 177), + ]; + let owned_frames = frames + .iter() + .map(|frame| { + let DynamicImage::ImageRgb8(rgb) = frame else { + panic!("test frame is not RGB8"); + }; + let mut bgr = rgb.as_raw().clone(); + for pixel in bgr.chunks_exact_mut(3) { + pixel.swap(0, 2); + } + OwnedRgbFrame { + width: rgb.width(), + height: rgb.height(), + data: bytes::Bytes::from(bgr), + channel_order: RgbChannelOrder::Bgr, + } + }) + .collect::>(); + let frame_refs = frames + .iter() + .map(|frame| { + let DynamicImage::ImageRgb8(rgb) = frame else { + panic!("test frame is not RGB8"); + }; + RgbFrameRef { + width: rgb.width(), + height: rgb.height(), + data: rgb.as_raw(), + } + }) + .collect::>(); + let mut bulk = processor + .preprocess_video_rgb_deferred(&frame_refs, &config) + .unwrap(); + + let (sender, receiver) = sync_channel(2); + let expected_frames = owned_frames.len(); + let producer = thread::spawn(move || { + for frame in owned_frames { + sender.send(Ok(frame)).unwrap(); + } + }); + let mut streamed = processor + .preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Video(VideoInput::RgbStream(DecodedRgbFrameStream::new( + expected_frames, + receiver, + ))), + output: OutputPreference::CompactAllowed, + config: &config, + }) + .unwrap(); + producer.join().unwrap(); + + let bulk_compact = bulk.encoder_input.deferred_normalized().unwrap(); + let streamed_compact = streamed.encoder_input.deferred_normalized().unwrap(); + let mut bulk_bf16 = vec![0; bulk_compact.len() * 2]; + let mut streamed_bf16 = vec![0; streamed_compact.len() * 2]; + bulk_compact.fill_bf16_le_bytes(&mut bulk_bf16).unwrap(); + streamed_compact + .fill_bf16_le_bytes(&mut streamed_bf16) + .unwrap(); + assert_eq!(streamed_bf16, bulk_bf16); + assert_eq!(streamed.feature_token_counts, bulk.feature_token_counts); + + bulk.materialize_encoder_input().unwrap(); + streamed.materialize_encoder_input().unwrap(); + for (index, (&actual, &expected)) in streamed + .encoder_input + .dense() + .unwrap() + .iter() + .zip(bulk.encoder_input.dense().unwrap().iter()) + .enumerate() + { + assert_eq!( + actual.to_bits(), + expected.to_bits(), + "streamed video path diverges at index {index}" + ); + } + } + + #[test] + fn test_streaming_bgr_video_matches_bulk_bicubic_with_padding() { + assert_bgr_streaming_matches_bulk(Some(3)); + } + + #[test] + fn test_streaming_bgr_video_matches_bulk_bilinear_with_padding() { + assert_bgr_streaming_matches_bulk(Some(2)); + } + #[test] fn test_qwen_vl_base_factor() { let processor = QwenVLProcessorBase::new(create_test_config()); @@ -1290,4 +2754,98 @@ mod tests { ); } } + + #[test] + fn test_preprocess_video_rgb_matches_dynamic_video_parallel_blocks() { + let processor = QwenVLProcessorBase::new(create_video_test_config()); + let config = PreProcessorConfig { + image_mean: Some(processor.default_mean().to_vec()), + image_std: Some(processor.default_std().to_vec()), + ..Default::default() + }; + let frames = vec![ + create_sized_pattern_frame(280, 280, 3), + create_sized_pattern_frame(280, 280, 101), + ]; + + let dynamic_result = processor.preprocess_video(&frames, &config).unwrap(); + let expected = dynamic_result + .encoder_input + .as_slice_memory_order() + .unwrap(); + + let video_frames = frames + .iter() + .map(|frame| { + let DynamicImage::ImageRgb8(rgb) = frame else { + panic!("test frame is not RGB8"); + }; + VideoFrameRgb { + width: rgb.width() as usize, + height: rgb.height() as usize, + data: Cow::Borrowed(rgb.as_raw()), + } + }) + .collect::>(); + let temporal_patch_size = processor.config.temporal_patch_size; + let padded_frames = frames.len().div_ceil(temporal_patch_size) * temporal_patch_size; + let (target_h, target_w) = processor + .smart_resize_video(frames.len(), 280, 280) + .unwrap(); + let (_grid_t, grid_h, grid_w) = + processor.calculate_grid_thw(target_h, target_w, padded_frames); + let patch_size = processor.config.patch_size; + let merge_size = processor.config.merge_size; + let merged_patch = merge_size * patch_size; + let pr_blocks = grid_h / merge_size; + let pc_blocks = grid_w / merge_size; + let n_blocks = pr_blocks * pc_blocks; + assert!( + n_blocks > 1, + "test must exercise multiple video patch blocks" + ); + let block_out = merge_size * merge_size * 3 * frames.len() * patch_size * patch_size; + let mean = processor.default_mean(); + let std = processor.default_std(); + let scale: [f32; 3] = std::array::from_fn(|c| 1.0 / (255.0 * std[c] as f32)); + let bias: [f32; 3] = std::array::from_fn(|c| -(mean[c] as f32) / (std[c] as f32)); + let lut: [[f32; 256]; 3] = + std::array::from_fn(|c| std::array::from_fn(|v| v as f32 * scale[c] + bias[c])); + + let mut actual = vec![0.0; expected.len()]; + let split_blocks = n_blocks / 2; + let split_at = split_blocks * block_out; + let (first, second) = actual.split_at_mut(split_at); + QwenVLProcessorBase::patchify_video_rgb_block_band( + &video_frames, + target_w, + patch_size, + merge_size, + merged_patch, + pc_blocks, + 0, + first, + &lut, + ); + QwenVLProcessorBase::patchify_video_rgb_block_band( + &video_frames, + target_w, + patch_size, + merge_size, + merged_patch, + pc_blocks, + split_blocks, + second, + &lut, + ); + + assert_eq!(actual.len(), expected.len()); + for (idx, (&got, &want)) in actual.iter().zip(expected.iter()).enumerate() { + assert_eq!( + got.to_bits(), + want.to_bits(), + "parallel RGB video patch value differs at index {idx}: got {got}, want {want}" + ); + } + } } diff --git a/crates/multimodal/src/vision/transforms.rs b/crates/multimodal/src/vision/transforms.rs index 06a8e091f..633d9bf1e 100644 --- a/crates/multimodal/src/vision/transforms.rs +++ b/crates/multimodal/src/vision/transforms.rs @@ -3,7 +3,7 @@ //! This module provides composable transforms that match HuggingFace image processor //! behavior, enabling pure Rust preprocessing without Python dependencies. -use std::cell::RefCell; +use std::{borrow::Cow, cell::RefCell}; use fast_image_resize::{ images::{Image as FirImage, ImageRef as FirImageRef}, @@ -39,18 +39,18 @@ pub type Result = std::result::Result; /// Extract RGB pixel data from a DynamicImage, avoiding a copy when already RGB8. /// Returns (width, height, raw_bytes) where raw_bytes is interleaved R,G,B,R,G,B,... -pub fn rgb_bytes(image: &DynamicImage) -> (usize, usize, std::borrow::Cow<'_, [u8]>) { +pub fn rgb_bytes(image: &DynamicImage) -> (usize, usize, Cow<'_, [u8]>) { match image { DynamicImage::ImageRgb8(rgb) => ( rgb.width() as usize, rgb.height() as usize, - std::borrow::Cow::Borrowed(rgb.as_raw()), + Cow::Borrowed(rgb.as_raw()), ), _ => { let rgb = image.to_rgb8(); let w = rgb.width() as usize; let h = rgb.height() as usize; - (w, h, std::borrow::Cow::Owned(rgb.into_raw())) + (w, h, Cow::Owned(rgb.into_raw())) } } } @@ -83,7 +83,7 @@ pub fn deinterleave_rgb_to_planes( } let chunk = pixels.div_ceil(nthreads); let (mut rr, mut gg, mut bb) = (r_plane, g_plane, b_plane); - std::thread::scope(|s| { + par_scope(|s| { let mut p0 = 0usize; while p0 < pixels { let n = chunk.min(pixels - p0); @@ -94,7 +94,7 @@ pub fn deinterleave_rgb_to_planes( gg = gt; bb = bt; let rgb_band = &rgb[p0 * 3..(p0 + n) * 3]; - s.spawn(move || deinterleave_contiguous(rgb_band, rb, gb, bbnd, scale, bias)); + s.spawn(move |_| deinterleave_contiguous(rgb_band, rb, gb, bbnd, scale, bias)); p0 += n; } }); @@ -326,14 +326,12 @@ fn fir_image_to_dynamic( // --------------------------------------------------------------------------- // Pillow-exact bicubic resize. // -// HuggingFace's (slow) `Qwen2VLImageProcessor` — which vLLM uses for Qwen2/3-VL -// — resizes via `PIL.Image.resize(size, BICUBIC)` on the uint8 image. The SIMD -// `fast_image_resize` path above is the same filter *family* (Catmull-Rom, -// a=-0.5) but diverges bit-wise on non-integer ratios (support scaling + -// fixed-point details), which the vision encoder amplifies into a large -// embedding shift vs vLLM. This routine replicates Pillow's `Resample.c` -// algorithm exactly (validated bit-for-bit against Pillow) so SMG's encoder -// inputs match HF/vLLM. +// Qwen image processors resize via `PIL.Image.resize(size, BICUBIC)` on the +// uint8 image. The SIMD `fast_image_resize` path above is the same filter +// *family* (Catmull-Rom, a=-0.5) but diverges bit-wise on non-integer ratios +// (support scaling + fixed-point details), which the vision encoder amplifies +// into a large embedding shift. This routine replicates Pillow's `Resample.c` +// algorithm exactly, validated against Pillow. const PIL_PRECISION_BITS: i64 = 32 - 8 - 2; const PIL_BICUBIC_SUPPORT: f64 = 2.0; @@ -416,23 +414,127 @@ fn pil_clip8(v: i64) -> u8 { } } -/// Number of threads to split a resample pass across. Each output row is an -/// independent fixed-point integer sum, so banding rows over threads yields -/// BIT-IDENTICAL output (no shared accumulation, inner sum order unchanged). -/// Small images run serial to avoid thread-spawn overhead. +type PilCoefficients = (Vec<(usize, usize)>, Vec>); + +pub(crate) struct PilBicubicRgbPlan { + in_w: usize, + in_h: usize, + out_w: usize, + out_h: usize, + horizontal: Option, + vertical: Option, +} + +impl PilBicubicRgbPlan { + pub(crate) fn new(in_w: u32, in_h: u32, out_w: u32, out_h: u32) -> Result { + if in_w == 0 || in_h == 0 || out_w == 0 || out_h == 0 { + return Err(TransformError::ShapeError( + "PIL bicubic resize dimensions must be non-zero".to_string(), + )); + } + let (in_w, in_h, out_w, out_h) = + (in_w as usize, in_h as usize, out_w as usize, out_h as usize); + Ok(Self { + in_w, + in_h, + out_w, + out_h, + horizontal: (in_w != out_w).then(|| pil_precompute_coeffs(in_w, out_w)), + vertical: (in_h != out_h).then(|| pil_precompute_coeffs(in_h, out_h)), + }) + } + + pub(crate) fn prepare_horizontal<'a>(&self, data: &'a [u8]) -> Result> { + let expected = self.in_w * self.in_h * 3; + if data.len() != expected { + return Err(TransformError::ShapeError(format!( + "PIL bicubic RGB source has {} bytes, expected {expected}", + data.len() + ))); + } + let Some((bounds, kernels)) = &self.horizontal else { + return Ok(Cow::Borrowed(data)); + }; + + let half = 1_i64 << (PIL_PRECISION_BITS - 1); + let row_out = self.out_w * 3; + let mut output = vec![0; self.in_h * row_out]; + let workers = par_threads(output.len(), self.in_h); + if workers <= 1 { + pil_h_band( + data, + bounds, + kernels, + half, + self.in_w, + self.out_w, + 3, + 0, + &mut output, + ); + } else { + let rows_per_task = self.in_h.div_ceil(workers); + par_scope(|scope| { + for (task, band) in output.chunks_mut(rows_per_task * row_out).enumerate() { + let first_row = task * rows_per_task; + scope.spawn(move |_| { + pil_h_band( + data, bounds, kernels, half, self.in_w, self.out_w, 3, first_row, band, + ); + }); + } + }); + } + Ok(Cow::Owned(output)) + } + + #[inline] + pub(crate) fn pixel(&self, horizontal: &[u8], x: usize, y: usize, channel: usize) -> u8 { + debug_assert!(x < self.out_w && y < self.out_h && channel < 3); + let Some((bounds, kernels)) = &self.vertical else { + return horizontal[(y * self.out_w + x) * 3 + channel]; + }; + let (source_y, rows) = bounds[y]; + let kernel = &kernels[y]; + let mut value = 1_i64 << (PIL_PRECISION_BITS - 1); + for row in 0..rows { + value += + horizontal[((source_y + row) * self.out_w + x) * 3 + channel] as i64 * kernel[row]; + } + pil_clip8(value) + } +} + +const PREPROCESS_PAR_MIN_BYTES: usize = 1 << 19; +const PREPROCESS_PAR_MIN_ROWS: usize = 32; +const PREPROCESS_PAR_MAX_THREADS: usize = 8; + +#[doc(hidden)] +pub fn preprocess_parallelism(output_bytes: usize, work_items: usize) -> usize { + par_threads(output_bytes, work_items) +} + +pub(crate) fn par_scope<'scope, OP, R>(op: OP) -> R +where + OP: FnOnce(&rayon::Scope<'scope>) -> R + Send, + R: Send, +{ + rayon::scope(op) +} + +/// Number of threads to split an elementwise or row-banded preprocessing pass +/// across. Each output row/element is independent, so banding work over threads +/// yields BIT-IDENTICAL output: no shared accumulation and no inner-loop order +/// changes. Small images run serial to avoid thread-spawn overhead. pub(crate) fn par_threads(out_bytes: usize, out_rows: usize) -> usize { - const PAR_MIN_BYTES: usize = 1 << 19; // ~512 KiB output; below this, serial - const MIN_ROWS_PER_THREAD: usize = 32; // keep enough work per thread - const MAX_THREADS: usize = 32; // spawning hundreds of threads costs more than it saves - if out_bytes < PAR_MIN_BYTES || out_rows < 2 * MIN_ROWS_PER_THREAD { + if out_bytes < PREPROCESS_PAR_MIN_BYTES || out_rows < PREPROCESS_PAR_MIN_ROWS.saturating_mul(2) + { return 1; } - let avail = std::thread::available_parallelism() - .map(|n| n.get()) - .unwrap_or(1); - (out_rows / MIN_ROWS_PER_THREAD) + let avail = rayon::current_num_threads(); + (out_rows / PREPROCESS_PAR_MIN_ROWS) .min(avail) - .clamp(1, MAX_THREADS) + .clamp(1, PREPROCESS_PAR_MAX_THREADS) } /// Process output rows `[oy0, oy0 + out_band.len()/row_out)` of the horizontal @@ -491,7 +593,7 @@ fn pil_resample_horizontal( ); } else { let chunk_rows = rows.div_ceil(nthreads); - std::thread::scope(|s| { + par_scope(|s| { let (b, k) = (&bounds, &kernels); let mut rest = out.as_mut_slice(); let mut oy0 = 0usize; @@ -500,7 +602,7 @@ fn pil_resample_horizontal( let (band, tail) = rest.split_at_mut(n * row_out); rest = tail; let start = oy0; - s.spawn(move || { + s.spawn(move |_| { pil_h_band(src, b, k, half, in_w, out_w, channels, start, band); }); oy0 += n; @@ -560,7 +662,7 @@ fn pil_resample_vertical( pil_v_band(src, &bounds, &kernels, half, width, channels, 0, &mut out); } else { let chunk_rows = out_h.div_ceil(nthreads); - std::thread::scope(|s| { + par_scope(|s| { let (b, k) = (&bounds, &kernels); let mut rest = out.as_mut_slice(); let mut oy0 = 0usize; @@ -569,7 +671,7 @@ fn pil_resample_vertical( let (band, tail) = rest.split_at_mut(n * row_out); rest = tail; let start = oy0; - s.spawn(move || pil_v_band(src, b, k, half, width, channels, start, band)); + s.spawn(move |_| pil_v_band(src, b, k, half, width, channels, start, band)); oy0 += n; } }); @@ -577,28 +679,26 @@ fn pil_resample_vertical( out } -/// Pillow-exact BICUBIC resize (RGB8). Horizontal pass then vertical pass with -/// an intermediate u8 buffer, matching `PIL.Image.resize(.., BICUBIC)`. +/// Pillow-exact BICUBIC resize (RGB8), matching +/// `PIL.Image.resize(.., BICUBIC)`. pub fn resize_bicubic_pil(image: &DynamicImage, out_w: u32, out_h: u32) -> DynamicImage { let rgb = image.to_rgb8(); let (in_w, in_h) = rgb.dimensions(); - let (in_w, in_h, out_w_u, out_h_u) = - (in_w as usize, in_h as usize, out_w as usize, out_h as usize); - let horiz = pil_resample_horizontal(rgb.as_raw(), in_h, in_w, out_w_u, 3); - let vert = pil_resample_vertical(&horiz, in_h, out_w_u, out_h_u, 3); + let output = resize_bicubic_pil_bytes(rgb.as_raw(), in_w, in_h, out_w, out_h); #[expect( clippy::expect_used, - reason = "vert is exactly out_w*out_h*3 bytes by construction" + reason = "output is exactly out_w*out_h*3 bytes by construction" )] - DynamicImage::ImageRgb8(RgbImage::from_raw(out_w, out_h, vert).expect("pil resize buffer size")) + DynamicImage::ImageRgb8( + RgbImage::from_raw(out_w, out_h, output).expect("pil resize buffer size"), + ) } /// PIL-exact bicubic resize over borrowed interleaved RGB bytes. /// /// Byte-for-byte equivalent of [`resize_bicubic_pil`] but for the raw-RGB video -/// frame path (`preprocess_video_rgb`), so default-bicubic video frames match -/// HF/vLLM the same way images do. Returns an `RgbImage` to drop straight into -/// the existing [`resize_rgb_bytes`] call sites. +/// frame path (`preprocess_video_rgb`). Returns an `RgbImage` to drop straight +/// into the existing [`resize_rgb_bytes`] call sites. pub fn resize_bicubic_pil_rgb( data: &[u8], width: u32, @@ -606,12 +706,7 @@ pub fn resize_bicubic_pil_rgb( out_w: u32, out_h: u32, ) -> Result { - let (in_w, in_h, out_w_u, out_h_u) = ( - width as usize, - height as usize, - out_w as usize, - out_h as usize, - ); + let (in_w, in_h) = (width as usize, height as usize); let expected = in_w.saturating_mul(in_h).saturating_mul(3); if data.len() != expected { return Err(TransformError::ShapeError(format!( @@ -619,15 +714,30 @@ pub fn resize_bicubic_pil_rgb( data.len() ))); } - let horiz = pil_resample_horizontal(data, in_h, in_w, out_w_u, 3); - let vert = pil_resample_vertical(&horiz, in_h, out_w_u, out_h_u, 3); - RgbImage::from_raw(out_w, out_h, vert).ok_or_else(|| { + let output = resize_bicubic_pil_bytes(data, width, height, out_w, out_h); + RgbImage::from_raw(out_w, out_h, output).ok_or_else(|| { TransformError::ShapeError(format!( "failed to build PIL bicubic RGB image for {out_w}x{out_h}" )) }) } +fn resize_bicubic_pil_bytes(data: &[u8], in_w: u32, in_h: u32, out_w: u32, out_h: u32) -> Vec { + let (in_w, in_h, out_w, out_h) = (in_w as usize, in_h as usize, out_w as usize, out_h as usize); + if in_w == out_w && in_h == out_h { + data.to_vec() + } else if in_w == out_w { + pil_resample_vertical(data, in_h, in_w, out_h, 3) + } else { + let horiz = pil_resample_horizontal(data, in_h, in_w, out_w, 3); + if in_h == out_h { + horiz + } else { + pil_resample_vertical(&horiz, in_h, out_w, out_h, 3) + } + } +} + /// Resize image preserving aspect ratio, fitting within max dimensions. pub fn resize_to_fit( image: &DynamicImage, @@ -876,10 +986,8 @@ mod tests { } /// The raw-RGB video resizer must be byte-for-byte identical to the - /// DynamicImage PIL-bicubic resizer used for images, so default-bicubic - /// video frames match HF/vLLM exactly — the same bit-identity guarantee - /// images get via the fingerprint tests. Guards the video resize path added - /// for HF/vLLM parity (`preprocess_video_rgb`). + /// DynamicImage PIL-bicubic resizer used for images. Guards the video resize + /// path used by `preprocess_video_rgb`. #[test] fn resize_bicubic_pil_rgb_matches_dynamic_path() { let (src_w, src_h) = (37u32, 23u32); // non-aligned source, non-trivial ratios @@ -907,6 +1015,32 @@ mod tests { ); } + #[test] + fn resize_bicubic_pil_rgb_skips_identity_axes_bit_exactly() { + let (src_w, src_h) = (31u32, 23u32); + let mut data = vec![0u8; src_w as usize * src_h as usize * 3]; + for (index, value) in data.iter_mut().enumerate() { + *value = (index as u8).wrapping_mul(37).wrapping_add(11); + } + + for (out_w, out_h) in [(src_w, 17), (19, src_h), (src_w, src_h)] { + let horizontal = + pil_resample_horizontal(&data, src_h as usize, src_w as usize, out_w as usize, 3); + let expected = pil_resample_vertical( + &horizontal, + src_h as usize, + out_w as usize, + out_h as usize, + 3, + ); + let actual = resize_bicubic_pil_rgb(&data, src_w, src_h, out_w, out_h) + .unwrap() + .into_raw(); + + assert_eq!(actual, expected, "identity-axis fast path changed pixels"); + } + } + /// `resize_bicubic_pil_rgb` rejects a buffer whose length doesn't match the /// declared dimensions rather than reading out of bounds. #[test] diff --git a/crates/multimodal/tests/preprocess_fingerprint.rs b/crates/multimodal/tests/preprocess_fingerprint.rs index 85c84b7f7..baf2e3ef6 100644 --- a/crates/multimodal/tests/preprocess_fingerprint.rs +++ b/crates/multimodal/tests/preprocess_fingerprint.rs @@ -54,7 +54,7 @@ const EXPECTED: &[u64] = &[0x391ca5deba1ff255, 0x5bde4728a72eba9d, 0x617d3e39f58 fn fingerprint(w: u32, h: u32) -> (u64, usize) { let proc = Qwen3VLProcessor::new(); let res = proc.preprocess(&[make(w, h)], &config()).unwrap(); - let flat = res.encoder_input_flat(); + let flat = res.encoder_input_flat().unwrap(); (fnv1a_f32(flat.as_ref()), flat.len()) } diff --git a/crates/multimodal/tests/vision_golden_tests.rs b/crates/multimodal/tests/vision_golden_tests.rs index 5baa57993..be0a55f46 100644 --- a/crates/multimodal/tests/vision_golden_tests.rs +++ b/crates/multimodal/tests/vision_golden_tests.rs @@ -153,7 +153,10 @@ fn run_golden_test(mode: &str, image_name: &str) { .preprocess(&[image], &config) .expect("Processing failed"); - let diff = max_diff(&golden, &result.encoder_input); + let diff = max_diff( + &golden, + result.encoder_input.dense().expect("dense encoder input"), + ); println!("{mode} - {image_name} image - Max difference: {diff:.6}"); println!("Golden shape: {:?}", golden.shape()); println!("Rust shape: {:?}", result.encoder_input.shape()); @@ -382,7 +385,7 @@ fn run_qwen2_vl_golden_test(image_name: &str) { ); // pixel_values is now already patchified: [total_patches, patch_features] - let rust_patches = result.encoder_input_flat(); + let rust_patches = result.encoder_input_flat().expect("dense encoder input"); let rust_shape = ( result.encoder_input.shape()[0], result.encoder_input.shape()[1], @@ -533,7 +536,7 @@ fn run_qwen3_vl_golden_test(image_name: &str) { ); // pixel_values is now already patchified: [total_patches, patch_features] - let rust_patches = result.encoder_input_flat(); + let rust_patches = result.encoder_input_flat().expect("dense encoder input"); let rust_shape = ( result.encoder_input.shape()[0], result.encoder_input.shape()[1], @@ -811,6 +814,8 @@ fn run_phi3_vision_golden_test(image_name: &str) { // Convert rust ArrayD to Array5 for comparison let rust_pixels = result .encoder_input + .dense() + .expect("dense encoder input") .clone() .into_dimensionality::() .expect("Failed to convert to Ix5"); @@ -1012,6 +1017,8 @@ fn run_phi4_vision_golden_test(image_name: &str) { // Compare pixel values let rust_pixels = result .encoder_input + .dense() + .expect("dense encoder input") .clone() .into_dimensionality::() .expect("Failed to convert to Ix5"); @@ -1223,7 +1230,7 @@ fn run_llama4_vision_golden_test(image_name: &str) { ); // Compare pixel values - let rust_pixels = result.encoder_input_flat(); + let rust_pixels = result.encoder_input_flat().expect("dense encoder input"); let num_golden_elements: usize = golden_shape.iter().product(); // Find the max difference for the actual tiles (not padding) @@ -1412,7 +1419,7 @@ fn run_pixtral_golden_test(image_name: &str) { ); // Compare pixel values - only compare the actual image region, not padding - let rust_pixels = result.encoder_input_flat(); + let rust_pixels = result.encoder_input_flat().expect("dense encoder input"); let golden_pixels_flat: Vec = golden_pixels.iter().copied().collect(); // Calculate indices for the actual image region (not padding) diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index 67b9caef2..78a96672c 100644 --- a/docs/reference/configuration.md +++ b/docs/reference/configuration.md @@ -942,10 +942,13 @@ smg \ These env-only variables tune how the router ships preprocessed multimodal tensors (image/video encoder inputs) to a TokenSpeed worker. They do not affect accuracy — the inline and shared-memory paths produce byte-identical tensors. +SHM handles include offsets; multi-item TokenSpeed encoder inputs may share one +packed segment while preserving the same byte-exact tensor payloads and reducing +per-tensor file lifecycle overhead. | Environment Variable | Default | Description | |---------------------|---------|-------------| -| `SMG_TOKENSPEED_MM_TENSOR_TRANSPORT` | `inline` | Transport for large MM tensors: `inline` (gRPC bytes), `shm` (always use `/dev/shm`), or `auto` (use `/dev/shm` only when the worker is *verified* to share it). In `auto`, the router compares the worker's advertised `/dev/shm` namespace token (`GetServerInfo`) to its own and uses SHM only on a match; otherwise it falls back to inline. No locality configuration is needed. | +| `SMG_TOKENSPEED_MM_TENSOR_TRANSPORT` | image/audio: `inline`; video: `auto` | Transport for large MM tensors: `inline` (gRPC bytes), `shm` (always use `/dev/shm`), or `auto` (use `/dev/shm` only when the worker is *verified* to share it). When unset, image/audio stay inline while video uses `auto` to avoid the high-throughput video gRPC byte-copy path on colocated workers without hurting image TTFT. In `auto`, the router compares the worker's advertised `/dev/shm` namespace token (`GetServerInfo`) to its own and uses SHM only on a match; otherwise it falls back to inline. No locality configuration is needed. | | `SMG_TOKENSPEED_MM_SHM_MIN_BYTES` | `65536` | Minimum tensor size (bytes) before the SHM path is used; smaller tensors stay inline. | | `SMG_LOG_MM_TIMING` | `false` | Log per-stage multimodal preprocessing/assembly timing at `INFO`. Accepts `1`/`true`/`yes`. | diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py index a2791c156..738b63f3d 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py @@ -13,6 +13,7 @@ import dataclasses import json import logging +import math import os import re import time @@ -445,7 +446,9 @@ async def GetModelInfo( if "model_dtype" in fields: response_kwargs["model_dtype"] = dtype if "multimodal_encoder_dtype" in fields: - response_kwargs["multimodal_encoder_dtype"] = dtype + response_kwargs["multimodal_encoder_dtype"] = self._multimodal_encoder_dtype( + model_config, self.scheduler_info + ) return tokenspeed_scheduler_pb2.GetModelInfoResponse(**response_kwargs) # ------------------------------------------------------------------ @@ -957,105 +960,145 @@ def _mm_inputs_from_itemized_proto( im_token_id = None video_token_id = None total_started = time.perf_counter() if LOG_MM_TIMING else None + deferred_shm_unlinks: set[str] = set() + preserved_encoder_shm_names: set[str] = set() + completed = False - for item_proto in mm_inputs.items: - item_started = time.perf_counter() if LOG_MM_TIMING else None - modality = self._modality_from_proto(item_proto.modality) - if not item_proto.HasField("encoder_input"): - raise ValueError("MultimodalItem must include encoder_input") - - feature_started = time.perf_counter() if LOG_MM_TIMING else None - feature = self._feature_from_proto(item_proto.encoder_input, cast_to=model_dtype) - feature_elapsed_ms = ( - (time.perf_counter() - feature_started) * 1000 - if feature_started is not None - else None - ) - if LOG_MM_TENSOR_DATA: - encoder_input = item_proto.encoder_input - payload = encoder_input.WhichOneof("payload") - inline_nbytes = len(encoder_input.inline) if payload == "inline" else None - logger.info( - "Multimodal encoder_input received: modality=%s proto_dtype=%s " - "shape=%s payload=%s inline_nbytes=%s feature_type=%s " - "torch_dtype=%s cast_to=%s", - modality, - encoder_input.dtype, - list(encoder_input.shape), - payload, - inline_nbytes, - type(feature).__name__, - feature.dtype, - model_dtype, + try: + for item_proto in mm_inputs.items: + if ( + item_proto.HasField("encoder_input") + and item_proto.encoder_input.WhichOneof("payload") == "shm" + ): + preserved_encoder_shm_names.add( + self._validated_shm_name(item_proto.encoder_input.shm.name) + ) + + for item_proto in mm_inputs.items: + item_started = time.perf_counter() if LOG_MM_TIMING else None + modality = self._modality_from_proto(item_proto.modality) + if not item_proto.HasField("encoder_input"): + raise ValueError("MultimodalItem must include encoder_input") + + feature_started = time.perf_counter() if LOG_MM_TIMING else None + feature = self._feature_from_proto( + item_proto.encoder_input, + cast_to=model_dtype, + deferred_unlink_names=deferred_shm_unlinks, ) - model_started = time.perf_counter() if LOG_MM_TIMING else None - model_specific_data = { - name: self._tensor_from_proto(tensor_data, cast_to=model_dtype) - for name, tensor_data in item_proto.model_specific_tensors.items() - } - model_elapsed_ms = ( - (time.perf_counter() - model_started) * 1000 if model_started is not None else None - ) - self._validate_item_tensor_consistency(modality, model_specific_data) + if isinstance(feature, ShmTensorHandle): + preserved_encoder_shm_names.add(feature.shm_name) + feature_elapsed_ms = ( + (time.perf_counter() - feature_started) * 1000 + if feature_started is not None + else None + ) + if LOG_MM_TENSOR_DATA: + encoder_input = item_proto.encoder_input + payload = encoder_input.WhichOneof("payload") + inline_nbytes = len(encoder_input.inline) if payload == "inline" else None + logger.info( + "Multimodal encoder_input received: modality=%s proto_dtype=%s " + "shape=%s payload=%s inline_nbytes=%s feature_type=%s " + "torch_dtype=%s cast_to=%s", + modality, + encoder_input.dtype, + list(encoder_input.shape), + payload, + inline_nbytes, + type(feature).__name__, + feature.dtype, + model_dtype, + ) + model_started = time.perf_counter() if LOG_MM_TIMING else None + model_specific_data = {} + for name, tensor_data in item_proto.model_specific_tensors.items(): + if tensor_data.WhichOneof("payload") == "shm": + model_shm_name = self._validated_shm_name(tensor_data.shm.name) + if model_shm_name in preserved_encoder_shm_names: + raise ValueError( + "model-specific tensors must not share SHM segments " + "with preserved encoder_input handles" + ) + model_specific_data[name] = self._tensor_from_proto( + tensor_data, + cast_to=model_dtype, + unlink_after_read=False, + deferred_unlink_names=deferred_shm_unlinks, + ) + model_elapsed_ms = ( + (time.perf_counter() - model_started) * 1000 + if model_started is not None + else None + ) + self._validate_item_tensor_consistency(modality, model_specific_data) - if not item_proto.placeholders: - raise ValueError("MultimodalItem carried no placeholders") - if any(p.length <= 0 for p in item_proto.placeholders): - raise ValueError("MultimodalItem.placeholders.length must be > 0") - offsets = [(p.offset, p.offset + p.length - 1) for p in item_proto.placeholders] - - content_hash = bytes(item_proto.content_hash) - mm_item = MultimodalDataItem( - modality=modality, - feature=feature, - model_specific_data=model_specific_data, - offsets=offsets, - hash=int.from_bytes(content_hash[:8], "little") if content_hash else None, - ) - mm_item.set_pad_value() - items.append(mm_item) + offsets, token_count, offset_ends, offset_prefix = ( + self._offsets_from_proto_placeholders(item_proto.placeholders) + ) - if LOG_MM_TIMING and item_started is not None: - encoder_input = item_proto.encoder_input - logger.info( - "mm_timing item_build_ms modality=%s elapsed=%.3f " - "feature_ms=%.3f model_specific_ms=%.3f payload=%s " - "proto_dtype=%s shape=%s feature_type=%s tensors=%s", - modality.name, - (time.perf_counter() - item_started) * 1000, - feature_elapsed_ms, - model_elapsed_ms, - encoder_input.WhichOneof("payload"), - encoder_input.dtype, - list(encoder_input.shape), - type(feature).__name__, - sorted(model_specific_data.keys()), + content_hash = bytes(item_proto.content_hash) + mm_item = MultimodalDataItem( + modality=modality, + feature=feature, + model_specific_data=model_specific_data, + offsets=offsets, + token_count=token_count, + hash=int.from_bytes(content_hash[:8], "little") if content_hash else None, + offset_ends=offset_ends, + offset_prefix=offset_prefix, ) + mm_item.set_pad_value() + items.append(mm_item) - if item_proto.HasField("placeholder_token_id"): - placeholder_token_id = int(item_proto.placeholder_token_id) - if modality == Modality.IMAGE: - im_token_id = self._merge_placeholder_token_id( - im_token_id, placeholder_token_id, modality - ) - elif modality == Modality.VIDEO: - video_token_id = self._merge_placeholder_token_id( - video_token_id, placeholder_token_id, modality + if LOG_MM_TIMING and item_started is not None: + encoder_input = item_proto.encoder_input + logger.info( + "mm_timing item_build_ms modality=%s elapsed=%.3f " + "feature_ms=%.3f model_specific_ms=%.3f payload=%s " + "proto_dtype=%s shape=%s feature_type=%s tensors=%s", + modality.name, + (time.perf_counter() - item_started) * 1000, + feature_elapsed_ms, + model_elapsed_ms, + encoder_input.WhichOneof("payload"), + encoder_input.dtype, + list(encoder_input.shape), + type(feature).__name__, + sorted(model_specific_data.keys()), ) - if not items: - raise ValueError("MultimodalInputs.items is empty") - if LOG_MM_TIMING and total_started is not None: - logger.info( - "mm_timing mm_inputs_build_ms items=%d elapsed=%.3f", - len(items), - (time.perf_counter() - total_started) * 1000, + if item_proto.HasField("placeholder_token_id"): + placeholder_token_id = int(item_proto.placeholder_token_id) + if modality == Modality.IMAGE: + im_token_id = self._merge_placeholder_token_id( + im_token_id, placeholder_token_id, modality + ) + elif modality == Modality.VIDEO: + video_token_id = self._merge_placeholder_token_id( + video_token_id, placeholder_token_id, modality + ) + + if not items: + raise ValueError("MultimodalInputs.items is empty") + if LOG_MM_TIMING and total_started is not None: + logger.info( + "mm_timing mm_inputs_build_ms items=%d elapsed=%.3f", + len(items), + (time.perf_counter() - total_started) * 1000, + ) + result = MultimodalInputs( + mm_items=items, + im_token_id=im_token_id, + video_token_id=video_token_id, + pad_values_ready=True, ) - return MultimodalInputs( - mm_items=items, - im_token_id=im_token_id, - video_token_id=video_token_id, - ) + completed = True + return result + finally: + if not completed: + deferred_shm_unlinks.update(preserved_encoder_shm_names) + self._unlink_deferred_shm(deferred_shm_unlinks) @staticmethod def _modality_from_proto(modality: int) -> Modality: @@ -1092,10 +1135,51 @@ def _validate_item_tensor_consistency( if modality == Modality.VIDEO and not has_video_grid: raise ValueError("VIDEO MultimodalItem must carry video_grid_thw") + @staticmethod + def _offsets_from_proto_placeholders( + placeholders, + ) -> tuple[list[tuple[int, int]], int, list[int] | None, list[int] | None]: + if len(placeholders) == 1: + placeholder = placeholders[0] + length = int(placeholder.length) + if length <= 0: + raise ValueError("MultimodalItem.placeholders.length must be > 0") + start = int(placeholder.offset) + end = start + length - 1 + return [(start, end)], length, [end], [0, length] + + offsets = [] + offset_ends = [] + offset_prefix = [0] + sorted_non_overlapping = True + prev_end = -1 + token_count = 0 + for placeholder in placeholders: + length = int(placeholder.length) + if length <= 0: + raise ValueError("MultimodalItem.placeholders.length must be > 0") + start = int(placeholder.offset) + end = start + length - 1 + if start <= prev_end: + sorted_non_overlapping = False + offsets.append((start, end)) + offset_ends.append(end) + token_count += length + offset_prefix.append(token_count) + prev_end = end + if not offsets: + raise ValueError("MultimodalItem carried no placeholders") + if not sorted_non_overlapping: + return offsets, token_count, None, None + return offsets, token_count, offset_ends, offset_prefix + @staticmethod def _tensor_from_proto( tensor_data: tokenspeed_scheduler_pb2.TensorData, cast_to: torch.dtype | None = None, + *, + unlink_after_read: bool = True, + deferred_unlink_names: set[str] | None = None, ): """Reconstruct a torch.Tensor from a proto TensorData. @@ -1103,11 +1187,24 @@ def _tensor_from_proto( copied so it never aliases the transient proto bytes. """ shape = list(tensor_data.shape) - raw = TokenSpeedSchedulerServicer._tensor_payload_bytes(tensor_data) + if tensor_data.dtype == "bfloat16": + expected = math.prod(shape) * np.dtype(np.uint16).itemsize + np_dtype = None + else: + try: + np_dtype = np.dtype(tensor_data.dtype) + except TypeError as exc: + raise ValueError(f"Unsupported TensorData dtype: {tensor_data.dtype!r}") from exc + expected = math.prod(shape) * np_dtype.itemsize + TokenSpeedSchedulerServicer._validate_shm_nbytes_before_read(tensor_data, shape, expected) + raw = TokenSpeedSchedulerServicer._tensor_payload_bytes( + tensor_data, + unlink_after_read=unlink_after_read, + deferred_unlink_names=deferred_unlink_names, + ) if tensor_data.dtype == "bfloat16": # numpy has no bfloat16 — read the raw bits as uint16, reinterpret. - expected = int(np.prod(shape, dtype=np.int64)) * np.dtype(np.uint16).itemsize if len(raw) != expected: raise ValueError( f"TensorData byte length mismatch for bfloat16 shape={shape}: " @@ -1117,23 +1214,38 @@ def _tensor_from_proto( torch.bfloat16 ) else: - dtype = np.dtype(tensor_data.dtype) - expected = int(np.prod(shape, dtype=np.int64)) * dtype.itemsize if len(raw) != expected: raise ValueError( f"TensorData byte length mismatch for dtype={tensor_data.dtype}, " f"shape={shape}: expected {expected}, got {len(raw)}" ) - t = torch.from_numpy(np.frombuffer(raw, dtype=dtype).reshape(shape)) + t = torch.from_numpy(np.frombuffer(raw, dtype=np_dtype).reshape(shape)) if cast_to is not None and t.dtype != cast_to and t.is_floating_point(): return t.to(cast_to) return t.clone() + @staticmethod + def _validate_shm_nbytes_before_read( + tensor_data: tokenspeed_scheduler_pb2.TensorData, + shape: list[int], + expected: int, + ) -> None: + if tensor_data.WhichOneof("payload") != "shm": + return + nbytes = int(tensor_data.shm.nbytes) + if nbytes != expected: + raise ValueError( + f"TensorData.shm byte length mismatch for dtype={tensor_data.dtype}, " + f"shape={shape}: expected {expected}, got {nbytes}" + ) + @staticmethod def _feature_from_proto( tensor_data: tokenspeed_scheduler_pb2.TensorData, cast_to: torch.dtype | None = None, + *, + deferred_unlink_names: set[str] | None = None, ) -> torch.Tensor | ShmTensorHandle: """Reconstruct a feature tensor, preserving SHM handles when possible. @@ -1146,35 +1258,54 @@ def _feature_from_proto( return TokenSpeedSchedulerServicer._tensor_from_proto(tensor_data, cast_to=cast_to) dtype = TokenSpeedSchedulerServicer._torch_dtype_from_proto(tensor_data.dtype) - if ( - cast_to is not None - and dtype != cast_to - and torch.is_floating_point(torch.empty((), dtype=dtype)) - ): - return TokenSpeedSchedulerServicer._tensor_from_proto(tensor_data, cast_to=cast_to) - - shm = tensor_data.shm - if shm.offset != 0: - return TokenSpeedSchedulerServicer._tensor_from_proto(tensor_data, cast_to=cast_to) + if cast_to is not None and dtype != cast_to: + return TokenSpeedSchedulerServicer._tensor_from_proto( + tensor_data, + cast_to=cast_to, + unlink_after_read=deferred_unlink_names is None, + deferred_unlink_names=deferred_unlink_names, + ) shape = tuple(int(dim) for dim in tensor_data.shape) - expected = int(np.prod(shape, dtype=np.int64)) * torch.empty((), dtype=dtype).element_size() - if int(shm.nbytes) != expected: + expected = math.prod(shape) * TokenSpeedSchedulerServicer._torch_dtype_size(dtype) + shm = tensor_data.shm + offset = int(shm.offset) + nbytes = int(shm.nbytes) + if offset < 0: + raise ValueError( + f"TensorData.shm offset must be non-negative for shape={list(shape)}: {offset}" + ) + if nbytes != expected: raise ValueError( f"TensorData.shm byte length mismatch for dtype={tensor_data.dtype}, " - f"shape={list(shape)}: expected {expected}, got {int(shm.nbytes)}" + f"shape={list(shape)}: expected {expected}, got {nbytes}" ) name = TokenSpeedSchedulerServicer._validated_shm_name(shm.name) - return ShmTensorHandle(shm_name=name, shape=shape, dtype=dtype) + return ShmTensorHandle( + shm_name=name, + shape=shape, + dtype=dtype, + offset=offset, + nbytes=nbytes, + ) @staticmethod - def _tensor_payload_bytes(tensor_data: tokenspeed_scheduler_pb2.TensorData) -> bytes: + def _tensor_payload_bytes( + tensor_data: tokenspeed_scheduler_pb2.TensorData, + *, + unlink_after_read: bool = True, + deferred_unlink_names: set[str] | None = None, + ) -> bytes: payload = tensor_data.WhichOneof("payload") if payload == "inline": return bytes(tensor_data.inline) if payload == "shm": - return TokenSpeedSchedulerServicer._tensor_payload_bytes_from_shm(tensor_data.shm) + return TokenSpeedSchedulerServicer._tensor_payload_bytes_from_shm( + tensor_data.shm, + unlink_after_read=unlink_after_read, + deferred_unlink_names=deferred_unlink_names, + ) if payload == "remote": raise ValueError("TensorData.remote payload is not implemented yet") raise ValueError("TensorData payload is required") @@ -1182,6 +1313,9 @@ def _tensor_payload_bytes(tensor_data: tokenspeed_scheduler_pb2.TensorData) -> b @staticmethod def _tensor_payload_bytes_from_shm( shm_handle: tokenspeed_scheduler_pb2.ShmHandle, + *, + unlink_after_read: bool = True, + deferred_unlink_names: set[str] | None = None, ) -> bytes: name = TokenSpeedSchedulerServicer._validated_shm_name(shm_handle.name) @@ -1193,19 +1327,28 @@ def _tensor_payload_bytes_from_shm( finally: if fd is not None: os.close(fd) - if fd is not None and UNLINK_MM_SHM_AFTER_READ: - try: - os.unlink(path) - except FileNotFoundError: - pass if len(raw) != int(shm_handle.nbytes): raise ValueError( f"TensorData.shm byte length mismatch for name={shm_handle.name!r}: " f"expected {int(shm_handle.nbytes)}, got {len(raw)}" ) + if unlink_after_read: + TokenSpeedSchedulerServicer._unlink_deferred_shm({name}) + elif deferred_unlink_names is not None and UNLINK_MM_SHM_AFTER_READ: + deferred_unlink_names.add(name) return raw + @staticmethod + def _unlink_deferred_shm(names: set[str]) -> None: + if not UNLINK_MM_SHM_AFTER_READ: + return + for name in names: + try: + os.unlink(os.path.join("/dev/shm", name)) + except FileNotFoundError: + pass + @staticmethod def _validated_shm_name(name: str) -> str: name = name.lstrip("/") @@ -1223,6 +1366,14 @@ def _torch_dtype_from_proto(dtype: str) -> torch.dtype: return torch.float32 raise ValueError(f"Unsupported TensorData dtype for SHM feature: {dtype!r}") + @staticmethod + def _torch_dtype_size(dtype: torch.dtype) -> int: + if dtype is torch.float32: + return 4 + if dtype is torch.float16 or dtype is torch.bfloat16: + return 2 + return torch.empty((), dtype=dtype).element_size() + @staticmethod def _torch_dtype_to_proto(dtype: torch.dtype | None) -> str: if dtype is torch.bfloat16: @@ -1233,6 +1384,12 @@ def _torch_dtype_to_proto(dtype: torch.dtype | None) -> str: return "float32" return "" + @staticmethod + def _multimodal_encoder_dtype(model_config: Any, scheduler_info: dict) -> str: + return scheduler_info.get("multimodal_encoder_dtype") or ( + TokenSpeedSchedulerServicer._torch_dtype_to_proto(getattr(model_config, "dtype", None)) + ) + def _generated_output_ids( self, output: dict, diff --git a/grpc_servicer/tests/test_tokenspeed_multimodal_shm.py b/grpc_servicer/tests/test_tokenspeed_multimodal_shm.py new file mode 100644 index 000000000..8512adf0c --- /dev/null +++ b/grpc_servicer/tests/test_tokenspeed_multimodal_shm.py @@ -0,0 +1,366 @@ +import os +from types import SimpleNamespace + +import pytest +from smg_grpc_proto.generated import tokenspeed_scheduler_pb2 + +torch = pytest.importorskip("torch") +shm_transport = pytest.importorskip("tokenspeed.runtime.multimodal.shm_transport") +servicer_module = pytest.importorskip("smg_grpc_servicer.tokenspeed.servicer") +ShmTensorHandle = shm_transport.ShmTensorHandle +TokenSpeedSchedulerServicer = servicer_module.TokenSpeedSchedulerServicer + + +def test_multimodal_encoder_dtype_prefers_loaded_vision_dtype(): + model_config = SimpleNamespace(dtype=torch.float32) + + assert ( + TokenSpeedSchedulerServicer._multimodal_encoder_dtype( + model_config, {"multimodal_encoder_dtype": "bfloat16"} + ) + == "bfloat16" + ) + + +def test_multimodal_encoder_dtype_falls_back_to_model_dtype(): + model_config = SimpleNamespace(dtype=torch.float16) + + assert TokenSpeedSchedulerServicer._multimodal_encoder_dtype(model_config, {}) == "float16" + + +def _require_writable_dev_shm(): + if not os.path.isdir("/dev/shm") or not os.access("/dev/shm", os.W_OK): + pytest.skip("/dev/shm is not available or writable") + + +def test_feature_from_proto_preserves_offset_shm_handle(): + tensor = tokenspeed_scheduler_pb2.TensorData( + shape=[3, 4], + dtype="float32", + shm=tokenspeed_scheduler_pb2.ShmHandle( + name="smg-tokenspeed-test", + offset=128, + nbytes=3 * 4 * 4, + owner_id="smg:test", + ), + ) + + feature = TokenSpeedSchedulerServicer._feature_from_proto(tensor) + + assert isinstance(feature, ShmTensorHandle) + assert feature.shm_name == "smg-tokenspeed-test" + assert feature.shape == (3, 4) + assert feature.dtype is torch.float32 + assert feature.offset == 128 + assert feature.nbytes == 3 * 4 * 4 + + +def test_feature_from_proto_rejects_offset_shm_length_mismatch(): + tensor = tokenspeed_scheduler_pb2.TensorData( + shape=[3, 4], + dtype="float32", + shm=tokenspeed_scheduler_pb2.ShmHandle( + name="smg-tokenspeed-test", + offset=128, + nbytes=4, + owner_id="smg:test", + ), + ) + + with pytest.raises(ValueError, match="byte length mismatch"): + TokenSpeedSchedulerServicer._feature_from_proto(tensor) + + +def test_feature_from_proto_cast_defers_shared_shm_unlink(monkeypatch): + _require_writable_dev_shm() + monkeypatch.setattr(servicer_module, "UNLINK_MM_SHM_AFTER_READ", True) + name = f"smg-tokenspeed-test-cast-{os.getpid()}" + path = os.path.join("/dev/shm", name) + raw = torch.tensor([1.0, 2.0, 3.0, 4.0], dtype=torch.float32).numpy().tobytes() + with open(path, "wb") as f: + f.write(raw) + + deferred_unlinks: set[str] = set() + first = tokenspeed_scheduler_pb2.TensorData( + shape=[2], + dtype="float32", + shm=tokenspeed_scheduler_pb2.ShmHandle( + name=name, + offset=0, + nbytes=2 * 4, + owner_id="smg:test", + ), + ) + second = tokenspeed_scheduler_pb2.TensorData( + shape=[2], + dtype="float32", + shm=tokenspeed_scheduler_pb2.ShmHandle( + name=name, + offset=2 * 4, + nbytes=2 * 4, + owner_id="smg:test", + ), + ) + + try: + first_feature = TokenSpeedSchedulerServicer._feature_from_proto( + first, + cast_to=torch.float16, + deferred_unlink_names=deferred_unlinks, + ) + assert first_feature.dtype == torch.float16 + assert os.path.exists(path) + + second_feature = TokenSpeedSchedulerServicer._feature_from_proto( + second, + cast_to=torch.float16, + deferred_unlink_names=deferred_unlinks, + ) + assert second_feature.dtype == torch.float16 + assert deferred_unlinks == {name} + + TokenSpeedSchedulerServicer._unlink_deferred_shm(deferred_unlinks) + assert not os.path.exists(path) + finally: + try: + os.unlink(path) + except FileNotFoundError: + pass + + +def test_mm_inputs_rejects_model_specific_reusing_preserved_encoder_shm(monkeypatch): + _require_writable_dev_shm() + monkeypatch.setattr(servicer_module, "UNLINK_MM_SHM_AFTER_READ", True) + name = f"smg-tokenspeed-test-preserved-{os.getpid()}" + path = os.path.join("/dev/shm", name) + try: + with open(path, "wb") as f: + f.write(torch.tensor([1.0], dtype=torch.float32).numpy().tobytes()) + + mm_inputs = tokenspeed_scheduler_pb2.MultimodalInputs( + items=[ + tokenspeed_scheduler_pb2.MultimodalItem( + modality=tokenspeed_scheduler_pb2.IMAGE, + content_hash=b"hash", + encoder_input=tokenspeed_scheduler_pb2.TensorData( + shape=[1], + dtype="float32", + shm=tokenspeed_scheduler_pb2.ShmHandle( + name=name, + offset=0, + nbytes=4, + owner_id="smg:test", + ), + ), + model_specific_tensors={ + "image_grid_thw": tokenspeed_scheduler_pb2.TensorData( + shape=[1], + dtype="uint32", + shm=tokenspeed_scheduler_pb2.ShmHandle( + name=name, + offset=0, + nbytes=4, + owner_id="smg:test", + ), + ), + }, + placeholders=[ + tokenspeed_scheduler_pb2.PlaceholderRange(offset=0, length=1), + ], + ), + ] + ) + servicer = object.__new__(TokenSpeedSchedulerServicer) + + with pytest.raises(ValueError, match="must not share SHM segments"): + servicer._mm_inputs_from_itemized_proto(mm_inputs) + assert not os.path.exists(path) + finally: + try: + os.unlink(path) + except FileNotFoundError: + pass + + +def test_mm_inputs_rejects_model_shm_reused_by_later_encoder(monkeypatch): + _require_writable_dev_shm() + monkeypatch.setattr(servicer_module, "UNLINK_MM_SHM_AFTER_READ", True) + + class FakeShmTensorHandle: + def __init__(self, *, shm_name, shape, dtype, offset, nbytes): + self.shm_name = shm_name + self.shape = shape + self.dtype = dtype + self.offset = offset + self.nbytes = nbytes + + monkeypatch.setattr(servicer_module, "ShmTensorHandle", FakeShmTensorHandle) + first_name = f"smg-tokenspeed-test-first-{os.getpid()}" + shared_name = f"smg-tokenspeed-test-later-{os.getpid()}" + paths = [os.path.join("/dev/shm", name) for name in (first_name, shared_name)] + try: + for path in paths: + with open(path, "wb") as f: + f.write(torch.tensor([1.0], dtype=torch.float32).numpy().tobytes()) + + def tensor(name, dtype="float32"): + return tokenspeed_scheduler_pb2.TensorData( + shape=[1], + dtype=dtype, + shm=tokenspeed_scheduler_pb2.ShmHandle( + name=name, + offset=0, + nbytes=4, + owner_id="smg:test", + ), + ) + + mm_inputs = tokenspeed_scheduler_pb2.MultimodalInputs( + items=[ + tokenspeed_scheduler_pb2.MultimodalItem( + modality=tokenspeed_scheduler_pb2.IMAGE, + encoder_input=tensor(first_name), + model_specific_tensors={"image_grid_thw": tensor(shared_name, "uint32")}, + placeholders=[ + tokenspeed_scheduler_pb2.PlaceholderRange(offset=0, length=1), + ], + ), + tokenspeed_scheduler_pb2.MultimodalItem( + modality=tokenspeed_scheduler_pb2.IMAGE, + encoder_input=tensor(shared_name), + placeholders=[ + tokenspeed_scheduler_pb2.PlaceholderRange(offset=1, length=1), + ], + ), + ] + ) + servicer = object.__new__(TokenSpeedSchedulerServicer) + + with pytest.raises(ValueError, match="must not share SHM segments"): + servicer._mm_inputs_from_itemized_proto(mm_inputs) + assert all(not os.path.exists(path) for path in paths) + finally: + for path in paths: + try: + os.unlink(path) + except FileNotFoundError: + pass + + +def test_mm_inputs_constructor_failure_unlinks_preserved_encoder_shm(monkeypatch): + _require_writable_dev_shm() + monkeypatch.setattr(servicer_module, "UNLINK_MM_SHM_AFTER_READ", True) + name = f"smg-tokenspeed-test-constructor-failure-{os.getpid()}" + path = os.path.join("/dev/shm", name) + + class FakeMultimodalDataItem: + def __init__(self, **_kwargs): + pass + + def set_pad_value(self): + pass + + def fail_multimodal_inputs(**_kwargs): + raise RuntimeError("constructor failed") + + monkeypatch.setattr(servicer_module, "MultimodalDataItem", FakeMultimodalDataItem) + monkeypatch.setattr(servicer_module, "MultimodalInputs", fail_multimodal_inputs) + + try: + with open(path, "wb") as f: + f.write(torch.tensor([1.0], dtype=torch.float32).numpy().tobytes()) + + mm_inputs = tokenspeed_scheduler_pb2.MultimodalInputs( + items=[ + tokenspeed_scheduler_pb2.MultimodalItem( + modality=tokenspeed_scheduler_pb2.IMAGE, + encoder_input=tokenspeed_scheduler_pb2.TensorData( + shape=[1], + dtype="float32", + shm=tokenspeed_scheduler_pb2.ShmHandle( + name=name, + offset=0, + nbytes=4, + owner_id="smg:test", + ), + ), + placeholders=[ + tokenspeed_scheduler_pb2.PlaceholderRange(offset=0, length=1), + ], + ), + ] + ) + servicer = object.__new__(TokenSpeedSchedulerServicer) + + with pytest.raises(RuntimeError, match="constructor failed"): + servicer._mm_inputs_from_itemized_proto(mm_inputs) + assert not os.path.exists(path) + finally: + try: + os.unlink(path) + except FileNotFoundError: + pass + + +def test_tensor_from_proto_rejects_shm_length_mismatch_before_read(): + tensor = tokenspeed_scheduler_pb2.TensorData( + shape=[3, 4], + dtype="float32", + shm=tokenspeed_scheduler_pb2.ShmHandle( + name="smg-tokenspeed-test-does-not-need-to-exist", + offset=0, + nbytes=4, + owner_id="smg:test", + ), + ) + + with pytest.raises(ValueError, match="byte length mismatch"): + TokenSpeedSchedulerServicer._tensor_from_proto(tensor) + + +def test_tensor_from_proto_normalizes_invalid_dtype_to_value_error(): + tensor = tokenspeed_scheduler_pb2.TensorData( + shape=[1], + dtype="not-a-real-dtype", + inline=b"", + ) + + with pytest.raises(ValueError, match="Unsupported TensorData dtype"): + TokenSpeedSchedulerServicer._tensor_from_proto(tensor) + + +def test_offsets_from_proto_placeholders_validates_and_builds_offsets_once(): + placeholders = [ + tokenspeed_scheduler_pb2.PlaceholderRange(offset=10, length=3), + tokenspeed_scheduler_pb2.PlaceholderRange(offset=20, length=1), + ] + + assert TokenSpeedSchedulerServicer._offsets_from_proto_placeholders(placeholders) == ( + [(10, 12), (20, 20)], + 4, + [12, 20], + [0, 3, 4], + ) + + +def test_offsets_from_proto_placeholders_single_placeholder_fast_path(): + placeholders = [ + tokenspeed_scheduler_pb2.PlaceholderRange(offset=10, length=3), + ] + + assert TokenSpeedSchedulerServicer._offsets_from_proto_placeholders(placeholders) == ( + [(10, 12)], + 3, + [12], + [0, 3], + ) + + +def test_offsets_from_proto_placeholders_rejects_empty_and_non_positive_lengths(): + with pytest.raises(ValueError, match="no placeholders"): + TokenSpeedSchedulerServicer._offsets_from_proto_placeholders([]) + + with pytest.raises(ValueError, match="length must be > 0"): + TokenSpeedSchedulerServicer._offsets_from_proto_placeholders( + [tokenspeed_scheduler_pb2.PlaceholderRange(offset=10, length=0)] + ) diff --git a/model_gateway/Cargo.toml b/model_gateway/Cargo.toml index 14f73fb44..5c1b4a71b 100644 --- a/model_gateway/Cargo.toml +++ b/model_gateway/Cargo.toml @@ -43,6 +43,7 @@ chrono.workspace = true dashmap.workspace = true futures.workspace = true http.workspace = true +libc.workspace = true parking_lot.workspace = true prost.workspace = true prost-types.workspace = true @@ -57,6 +58,7 @@ tracing-subscriber.workspace = true # Workspace dependencies (with extra features) axum = { workspace = true, features = ["macros", "multipart", "ws", "tracing"] } bytemuck = { workspace = true, features = ["derive"] } +memmap2 = { workspace = true } reqwest = { workspace = true, features = ["stream", "blocking", "json", "rustls", "multipart"] } serde = { workspace = true, features = ["derive"] } tokio = { workspace = true, features = ["full"] } @@ -83,6 +85,7 @@ bincode = "1.3" clap = { version = "4", features = ["derive", "env"] } axum-server = { version = "0.8.0", default-features = false, features = ["tls-rustls-no-provider"] } ndarray = "0.17" +rayon = "1.12" tower = { version = "0.5", features = ["full"] } tower-http = { version = "0.7", features = ["trace", "compression-gzip", "cors", "timeout", "limit", "request-id", "util"] } serde_json = { version = "1.0", default-features = false, features = ["std", "preserve_order"] } diff --git a/model_gateway/src/routers/grpc/multimodal.rs b/model_gateway/src/routers/grpc/multimodal.rs index 2beded605..e788fd832 100644 --- a/model_gateway/src/routers/grpc/multimodal.rs +++ b/model_gateway/src/routers/grpc/multimodal.rs @@ -9,25 +9,30 @@ //! functions differ because they work with different input types (`ChatMessage` vs //! `InputMessage`). +mod assembly; + use std::{ collections::HashMap, - io::Write, - mem::size_of, path::Path, sync::{Arc, OnceLock}, time::Instant, }; use anyhow::{Context, Result}; +pub(crate) use assembly::assemble_multimodal_data; +#[cfg(test)] +use assembly::*; use dashmap::DashMap; +#[cfg(test)] +use llm_multimodal::EncoderInput; use llm_multimodal::{ AsyncMultiModalTracker, FieldLayout, ImageDetail, ImageFrame, MediaConnector, - MediaConnectorConfig, MediaContentPart, Modality, ModelMetadata, ModelRegistry, - ModelSpecificValue, PlaceholderRange, PreProcessorConfig, PreprocessedEncoderInputs, - PromptReplacement, TrackedMedia, TrackerOutput, VideoClip, VisionProcessorRegistry, + MediaConnectorConfig, MediaContentPart, Modality, ModalityPreProcessor, + ModalityProcessorRegistry, ModelMetadata, ModelRegistry, MultimodalRuntime, OutputPreference, + PlaceholderRange, PreProcessorConfig, PreprocessRequest, PreprocessedEncoderInputs, + PromptReplacement, TrackedMedia, TrackerOutput, VideoClip, VideoInput, VisionInput, }; use llm_tokenizer::TokenizerTrait; -use ndarray::{ArrayD, Axis, Slice}; use openai_protocol::{ chat::{ChatMessage, MessageContent}, common::ContentPart, @@ -35,19 +40,6 @@ use openai_protocol::{ }; use tracing::{debug, info, warn}; -use crate::routers::grpc::{ - client::GrpcClient, - context::WorkerSelection, - proto_wrapper::{ - cleanup_tokenspeed_items_encoder_shm, tokenspeed_mm_shm_min_bytes, - tokenspeed_mm_tensor_transport_mode, tokenspeed_shm_dev_writable, - write_tokenspeed_shm_with, SglangMultimodalData, TensorBytes, TokenSpeedModality, - TokenSpeedMultimodalData, TokenSpeedMultimodalItem, TokenSpeedTensor, TrtllmMultimodalData, - VllmMultimodalData, - }, - MultimodalData, -}; - /// Cached model configuration files loaded from the tokenizer directory. #[derive(Debug, Clone)] pub(crate) struct MultimodalModelConfig { @@ -69,9 +61,12 @@ pub struct MultimodalConfigRegistry { } fn log_mm_timing_enabled() -> bool { - std::env::var("SMG_LOG_MM_TIMING") - .map(|value| matches!(value.to_ascii_lowercase().as_str(), "1" | "true" | "yes")) - .unwrap_or(false) + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var("SMG_LOG_MM_TIMING") + .map(|value| matches!(value.to_ascii_lowercase().as_str(), "1" | "true" | "yes")) + .unwrap_or(false) + }) } impl MultimodalConfigRegistry { @@ -240,8 +235,9 @@ pub(crate) fn load_video_preprocessor_config(base_dir: &Path) -> Option, pub media_connector: Arc, - pub vision_processor_registry: Arc, + pub processor_registry: Arc, pub model_registry: Arc, /// Shared reference to the app-level multimodal config cache. pub config_registry: Arc, @@ -255,12 +251,19 @@ impl MultimodalComponents { .timeout(std::time::Duration::from_secs(30)) .build() .context("Failed to create reqwest client")?; - let media_connector = MediaConnector::new(client, MediaConnectorConfig::default()) - .context("Failed to create MediaConnector")?; + let runtime = + Arc::new(MultimodalRuntime::new().context("Failed to create multimodal runtime")?); + let media_connector = MediaConnector::new_with_runtime( + client, + MediaConnectorConfig::default(), + runtime.clone(), + ) + .context("Failed to create MediaConnector")?; Ok(Self { + runtime, media_connector: Arc::new(media_connector), - vision_processor_registry: Arc::new(VisionProcessorRegistry::with_defaults()), + processor_registry: Arc::new(ModalityProcessorRegistry::with_defaults()), model_registry: Arc::new(ModelRegistry::default()), config_registry, }) @@ -272,40 +275,113 @@ pub(crate) struct MultimodalOutput { /// Token IDs with placeholder tokens expanded to the correct count per media item. pub expanded_token_ids: Vec, /// Lightweight intermediate holding preprocessing results. - /// Assembled into backend-specific `MultimodalData` in request_building. + /// Assembled into backend-specific data in request_building. pub intermediate: MultimodalIntermediate, } /// Lightweight intermediate from the preparation stage. /// /// Holds all preprocessing results without serializing tensors to bytes. -/// The assembly stage converts this into a backend-specific [`MultimodalData`] -/// variant once the target backend is known (after worker selection). -#[derive(Debug)] -pub(crate) enum MultimodalIntermediate { - Precomputed(PrecomputedMultimodalIntermediate), -} - +/// The assembly stage converts this into backend-specific data +/// once the target backend is known (after worker selection). #[derive(Debug)] -pub(crate) struct PrecomputedMultimodalIntermediate { - /// Active modality for this preprocessed payload. - pub modality: Modality, +pub(crate) struct MultimodalIntermediate { + /// Fetched media in exactly one active modality. + pub media: PreparedMedia, /// Preprocessed encoder input and model-specific tensors (not yet serialized). - pub preprocessed: PreprocessedEncoderInputs, - /// Raw image frames (bytes + blake3 hashes). - pub images: Vec>, - /// Raw video clips (bytes + blake3 hashes + sampled frames). - pub videos: Vec>, - /// Full structural placeholder ranges (offset, length). - pub placeholders: Vec, - /// Patch-only placeholder offsets for sglang. - pub patch_offsets: Option>, + pub preprocessed: Arc, + /// Full structural token ranges occupied by each media item. + pub structural_ranges: Vec, + /// Encoder-feature token ranges inside the structural ranges, when known. + pub feature_ranges: Option>, /// Placeholder token ID from model config for the active modality. pub placeholder_token_id: Option, /// Per-tensor field layout classification from the model spec. pub field_layouts: HashMap, - /// Tensor keys that should remain on CPU (vLLM `keep_on_cpu` hint). - pub keep_on_cpu_keys: Vec, + /// Model-declared tensor keys that should remain CPU-resident. + pub cpu_resident_tensor_keys: Vec, +} + +/// Prepared media carried between preprocessing and backend assembly. +/// +/// Adding audio or another modality requires explicit handling in every +/// backend adapter, while preventing contradictory modality and payload fields. +#[derive(Debug, Clone)] +pub(crate) enum PreparedMedia { + Images(Vec>), + Videos(Vec>), +} + +impl PreparedMedia { + pub(crate) fn modality(&self) -> Modality { + match self { + Self::Images(_) => Modality::Image, + Self::Videos(_) => Modality::Video, + } + } + + pub(crate) fn item_count(&self) -> usize { + match self { + Self::Images(images) => images.len(), + Self::Videos(videos) => videos.len(), + } + } + + pub(crate) fn content_hash(&self, item_index: usize) -> Option<&str> { + match self { + Self::Images(images) => images.get(item_index).map(|image| image.hash.as_str()), + Self::Videos(videos) => videos.get(item_index).map(|video| video.hash.as_str()), + } + } +} + +fn prepared_media_from_tracker(output: TrackerOutput) -> Result { + let mut prepared = None; + + for (modality, items) in output.data { + if items.is_empty() { + continue; + } + + let media = match modality { + Modality::Image => PreparedMedia::Images( + items + .into_iter() + .map(|item| match item { + TrackedMedia::Image(image) => Ok(image), + _ => Err(anyhow::anyhow!( + "Tracker returned non-image media under the image modality" + )), + }) + .collect::>>()?, + ), + Modality::Video => PreparedMedia::Videos( + items + .into_iter() + .map(|item| match item { + TrackedMedia::Video(video) => Ok(video), + _ => Err(anyhow::anyhow!( + "Tracker returned non-video media under the video modality" + )), + }) + .collect::>>()?, + ), + Modality::Audio | Modality::ImageEmbeds => { + return Err(anyhow::anyhow!( + "Prepared media does not support {modality} inputs yet" + )); + } + }; + + if prepared.replace(media).is_some() { + return Err(anyhow::anyhow!( + "Mixed multimodal requests are not supported yet" + )); + } + } + + prepared + .ok_or_else(|| anyhow::anyhow!("No media was successfully fetched for multimodal request")) } /// Resolve the placeholder token string for a multimodal model. @@ -559,8 +635,8 @@ async fn process_multimodal_parts( tokenizer_source: &str, ) -> Result { let log_timing = log_mm_timing_enabled(); - let total_started = Instant::now(); - let media_started = Instant::now(); + let total_started = log_timing.then(Instant::now); + let media_started = log_timing.then(Instant::now); let mut tracker = AsyncMultiModalTracker::new(components.media_connector.clone()); for part in content_parts { @@ -569,85 +645,63 @@ async fn process_multimodal_parts( .map_err(|e| anyhow::anyhow!("Failed to push content part: {e}"))?; } - let tracker_output: TrackerOutput = tracker - .finalize() - .await - .map_err(|e| anyhow::anyhow!("Failed to finalize multimodal tracker: {e}"))?; - - let images: Vec> = tracker_output - .data - .get(&Modality::Image) - .map(|media_vec| { - media_vec - .iter() - .filter_map(|m| match m { - TrackedMedia::Image(frame) => Some(frame.clone()), - _ => None, - }) - .collect() - }) - .unwrap_or_default(); - - let videos: Vec> = tracker_output - .data - .get(&Modality::Video) - .map(|media_vec| { - media_vec - .iter() - .filter_map(|m| match m { - TrackedMedia::Video(clip) => Some(clip.clone()), - _ => None, - }) - .collect() - }) - .unwrap_or_default(); + let media_future = async move { + let tracker_output: TrackerOutput = tracker + .finalize() + .await + .map_err(|e| anyhow::anyhow!("Failed to finalize multimodal tracker: {e}"))?; + Ok::<(TrackerOutput, f64), anyhow::Error>(( + tracker_output, + media_started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(), + )) + }; - let media_elapsed_ms = media_started.elapsed().as_secs_f64() * 1000.0; - let modality = match (images.is_empty(), videos.is_empty()) { - (false, true) => Modality::Image, - (true, false) => Modality::Video, - (false, false) => { - return Err(anyhow::anyhow!( - "Mixed image and video multimodal requests are not supported yet" - )); - } - (true, true) => { - return Err(anyhow::anyhow!( - "No media was successfully fetched for multimodal request" - )); - } + let config_future = async { + let config_started = log_timing.then(Instant::now); + let model_config = components + .config_registry + .get_or_load(tokenizer_id, tokenizer_source) + .await?; + Ok::<(Arc, f64), anyhow::Error>(( + model_config, + config_started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(), + )) }; - if modality == Modality::Video && videos.len() != 1 { + let (tracker_result, config_result) = tokio::join!(media_future, config_future); + let (tracker_output, media_elapsed_ms) = tracker_result?; + let media = prepared_media_from_tracker(tracker_output)?; + let modality = media.modality(); + + if matches!(&media, PreparedMedia::Videos(videos) if videos.len() != 1) { return Err(anyhow::anyhow!( "Exactly one video is supported per request for the initial video path" )); } - match modality { - Modality::Image => { + match &media { + PreparedMedia::Images(images) => { debug!( image_count = images.len(), item_sizes = ?images.iter().map(|f| (f.image.width(), f.image.height())).collect::>(), "Fetched images for multimodal processing" ); } - Modality::Video => { + PreparedMedia::Videos(videos) => { debug!( video_count = videos.len(), - frame_count = videos.first().map_or(0, |v| v.frames.len()), + frame_count = videos.first().map_or(0, |v| v.frame_count()), "Fetched video for multimodal processing" ); } - _ => {} } // Step 2: Resolve model spec and preprocess media. - let config_started = Instant::now(); - let model_config = components - .config_registry - .get_or_load(tokenizer_id, tokenizer_source) - .await?; + let (model_config, config_elapsed_ms) = config_result?; let model_type = model_config .config .get("model_type") @@ -661,11 +715,9 @@ async fn process_multimodal_parts( .model_registry .lookup(&metadata) .ok_or_else(|| anyhow::anyhow!("Multimodal not supported for model: {model_id}"))?; - let config_elapsed_ms = config_started.elapsed().as_secs_f64() * 1000.0; - // Run CPU-intensive vision preprocessing on a blocking thread pool so it - // doesn't block the tokio async runtime under concurrent load. - // TODO: consider making the thread pool size configurable. + // Run CPU-intensive modality preprocessing on the owned blocking pool so + // it does not block the Tokio runtime under concurrent load. let pp_config = match modality { Modality::Video => model_config .video_preprocessor_config @@ -673,80 +725,58 @@ async fn process_multimodal_parts( .unwrap_or_else(|| model_config.preprocessor_config.clone()), _ => model_config.preprocessor_config.clone(), }; - let registry = components.vision_processor_registry.clone(); + let registry = components.processor_registry.clone(); let model_id_owned = model_id.to_string(); let model_type_owned = model_type.map(String::from); - let images_for_preprocess = images.clone(); // cheap Arc refcount bumps - let videos_for_preprocess = videos.clone(); // cheap Arc refcount bumps - - let preprocess_started = Instant::now(); - let preprocessed: PreprocessedEncoderInputs = tokio::task::spawn_blocking(move || { - let processor = registry - .find(&model_id_owned, model_type_owned.as_deref()) - .ok_or_else(|| { - anyhow::anyhow!("No vision processor found for model: {model_id_owned}") - })?; - match modality { - Modality::Image => { - // Extract DynamicImages inside the blocking closure so the expensive - // clone happens off the tokio async runtime. - let raw_images: Vec = images_for_preprocess - .iter() - .map(|f| f.image.clone()) - .collect(); - processor - .preprocess(&raw_images, &pp_config) - .map_err(|e| anyhow::anyhow!("Image preprocessing failed: {e}")) - } - Modality::Video => { - let video = videos_for_preprocess - .first() - .ok_or_else(|| anyhow::anyhow!("No video available for preprocessing"))?; - - if !video.frames().is_empty() { - return processor - .preprocess_video(video.frames(), &pp_config) - .map_err(|e| anyhow::anyhow!("Video preprocessing failed: {e}")); - } - - if let Some(rgb_video) = video.rgb_video() { - match rgb_video.frame_refs() { - Ok(frame_refs) => { - match processor.preprocess_video_rgb(&frame_refs, &pp_config) { - Ok(preprocessed) => return Ok(preprocessed), - Err(error) => { - warn!( - error = %error, - "RGB video preprocessing fast path failed; falling back to materialized frames" - ); - } - } - } - Err(error) => { - warn!( - error = %error, - "RGB video frame refs are invalid; falling back to materialized frames" - ); - } - } - } + let preprocess_started = log_timing.then(Instant::now); + let model_type_for_preprocess = model_type_owned.clone(); + let media_for_preprocess = media.clone(); // cheap Arc refcount bumps + let runtime = components.runtime.clone(); + let preprocess_task = tokio::task::spawn_blocking(move || { + runtime.run_cpu(|| { + preprocess_media( + ®istry, + &model_id_owned, + model_type_for_preprocess.as_deref(), + &media_for_preprocess, + &pp_config, + ) + }) + }); - let frames = video - .materialized_frames() - .map_err(|e| anyhow::anyhow!("Video frame materialization failed: {e}"))?; - processor - .preprocess_video(&frames, &pp_config) - .map_err(|e| anyhow::anyhow!("Video preprocessing failed: {e}")) + let placeholder_ids_result: Result<(Option, Option)> = (|| { + // These IDs depend only on the model config/tokenizer, so resolve them + // while CPU vision preprocessing is running on the blocking pool. + let placeholder_token = spec + .placeholder_token_for(&metadata, modality) + .map_err(|e| anyhow::anyhow!("Failed to get placeholder token: {e}"))?; + let search_token_id = tokenizer.token_to_id(&placeholder_token); + let placeholder_token_id: Option = match spec + .placeholder_token_id_for(&metadata, modality) + { + Ok(id) => Some(id as u32), + Err(e) => { + warn!( + error = %e, + ?search_token_id, + "Failed to resolve placeholder_token_id from config, falling back to tokenizer lookup" + ); + search_token_id } - _ => Err(anyhow::anyhow!( - "Unsupported modality for preprocessing: {modality}" - )), - } - }) - .await - .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))??; - let preprocess_elapsed_ms = preprocess_started.elapsed().as_secs_f64() * 1000.0; + }; + Ok((search_token_id, placeholder_token_id)) + })(); + + let preprocessed = Arc::new( + preprocess_task + .await + .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}")) + .and_then(|inner| inner)?, + ); + let preprocess_elapsed_ms = preprocess_started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(); debug!( ?modality, @@ -756,7 +786,7 @@ async fn process_multimodal_parts( ); // Step 3: Compute prompt replacements and expand tokens. - let expansion_started = Instant::now(); + let expansion_started = log_timing.then(Instant::now); let prompt_replacements = spec .prompt_replacements_for(&metadata, &preprocessed, modality) .map_err(|e| anyhow::anyhow!("Failed to compute prompt replacements: {e}"))?; @@ -764,22 +794,7 @@ async fn process_multimodal_parts( // Two token IDs may differ for the same placeholder: // - search_token_id: what the tokenizer actually emits (e.g. 200090 for "<|image|>") // - placeholder_token_id: what the model config declares (e.g. image_token_id/video_token_id) - let placeholder_token = spec - .placeholder_token_for(&metadata, modality) - .map_err(|e| anyhow::anyhow!("Failed to get placeholder token: {e}"))?; - let search_token_id = tokenizer.token_to_id(&placeholder_token); - let placeholder_token_id: Option = match spec.placeholder_token_id_for(&metadata, modality) - { - Ok(id) => Some(id as u32), - Err(e) => { - warn!( - error = %e, - ?search_token_id, - "Failed to resolve placeholder_token_id from config, falling back to tokenizer lookup" - ); - search_token_id - } - }; + let (search_token_id, placeholder_token_id) = placeholder_ids_result?; let expanded = expand_tokens( &token_ids, @@ -791,40 +806,46 @@ async fn process_multimodal_parts( debug!( original_len = token_ids.len(), expanded_len = expanded.token_ids.len(), - placeholder_count = expanded.placeholders.len(), + placeholder_count = expanded.structural_ranges.len(), ?search_token_id, ?placeholder_token_id, "Token expansion complete" ); - let expansion_elapsed_ms = expansion_started.elapsed().as_secs_f64() * 1000.0; - let image_count = images.len(); - let video_count = videos.len(); - let video_frame_count = videos.first().map_or(0, |video| { - if video.frames().is_empty() { - video - .rgb_video() - .map_or(0, |rgb_video| rgb_video.frames.len()) - } else { - video.frames().len() - } + let expansion_elapsed_ms = expansion_started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(); + let timing_counts = log_timing.then(|| { + let (image_count, video_count, video_frame_count) = match &media { + PreparedMedia::Images(images) => (images.len(), 0, 0), + PreparedMedia::Videos(videos) => ( + 0, + videos.len(), + videos.first().map_or(0, |video| video.frame_count()), + ), + }; + ( + image_count, + video_count, + video_frame_count, + token_ids.len(), + expanded.token_ids.len(), + ) }); - let original_tokens = token_ids.len(); - let expanded_tokens = expanded.token_ids.len(); // Step 4: Build lightweight intermediate (defers tensor serialization to assembly) - let intermediate = MultimodalIntermediate::Precomputed(PrecomputedMultimodalIntermediate { - modality, + let intermediate = MultimodalIntermediate { + media, preprocessed, - images, - videos, - placeholders: expanded.placeholders, - patch_offsets: expanded.patch_offsets, + structural_ranges: expanded.structural_ranges, + feature_ranges: expanded.feature_ranges, placeholder_token_id, field_layouts: spec.field_layouts(), - keep_on_cpu_keys: spec.keep_on_cpu_keys(), - }); + cpu_resident_tensor_keys: spec.cpu_resident_tensor_keys(), + }; - if log_timing { + if let Some((image_count, video_count, video_frame_count, original_tokens, expanded_tokens)) = + timing_counts + { info!( modality = ?modality, image_count, @@ -834,7 +855,9 @@ async fn process_multimodal_parts( config_lookup_ms = config_elapsed_ms, preprocess_ms = preprocess_elapsed_ms, token_expand_ms = expansion_elapsed_ms, - total_ms = total_started.elapsed().as_secs_f64() * 1000.0, + total_ms = total_started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(), original_tokens, expanded_tokens, "smg_mm_timing process_multimodal_parts" @@ -847,24 +870,121 @@ async fn process_multimodal_parts( }) } -/// Output of token expansion, containing both full structural and patch-only ranges. +fn preprocess_media( + registry: &ModalityProcessorRegistry, + model_id: &str, + model_type: Option<&str>, + media: &PreparedMedia, + config: &PreProcessorConfig, +) -> Result { + let processor = registry + .find(model_id, model_type) + .ok_or_else(|| anyhow::anyhow!("No modality processor found for model: {model_id}"))?; + + match media { + PreparedMedia::Images(images) => { + let raw_images: Vec<&image::DynamicImage> = + images.iter().map(|frame| &frame.image).collect(); + processor + .preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Images(&raw_images), + output: if images.len() == 1 { + OutputPreference::CompactAllowed + } else { + OutputPreference::Materialized + }, + config, + }) + .map_err(|error| anyhow::anyhow!("Image preprocessing failed: {error}")) + } + PreparedMedia::Videos(videos) => preprocess_video(processor, videos, config), + } +} + +fn preprocess_video( + processor: &dyn ModalityPreProcessor, + videos: &[Arc], + config: &PreProcessorConfig, +) -> Result { + let video = videos + .first() + .ok_or_else(|| anyhow::anyhow!("No video available for preprocessing"))?; + + if let Some(stream) = video + .take_rgb_stream() + .map_err(|error| anyhow::anyhow!("Video frame stream unavailable: {error}"))? + { + return processor + .preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Video(VideoInput::RgbStream(stream)), + output: OutputPreference::CompactAllowed, + config, + }) + .map_err(|error| anyhow::anyhow!("Video stream preprocessing failed: {error}")); + } + + if !video.frames().is_empty() { + return processor + .preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Video(VideoInput::Frames(video.frames())), + output: OutputPreference::Materialized, + config, + }) + .map_err(|error| anyhow::anyhow!("Video preprocessing failed: {error}")); + } + + if let Some(rgb_video) = video.rgb_video() { + match rgb_video.frame_refs() { + Ok(frame_refs) => match processor.preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Video(VideoInput::Rgb(&frame_refs)), + output: OutputPreference::CompactAllowed, + config, + }) { + Ok(preprocessed) => return Ok(preprocessed), + Err(error) => { + warn!( + error = %error, + "RGB video preprocessing fast path failed; falling back to materialized frames" + ); + } + }, + Err(error) => { + warn!( + error = %error, + "RGB video frame refs are invalid; falling back to materialized frames" + ); + } + } + } + + let frames = video + .materialized_frames() + .map_err(|error| anyhow::anyhow!("Video frame materialization failed: {error}"))?; + processor + .preprocess_input(PreprocessRequest::Vision { + input: VisionInput::Video(VideoInput::Frames(&frames)), + output: OutputPreference::Materialized, + config, + }) + .map_err(|error| anyhow::anyhow!("Video preprocessing failed: {error}")) +} + +/// Output of token expansion with structural and encoder-feature ranges. struct ExpandedTokens { /// The expanded token ID sequence. token_ids: Vec, - /// Full structural placeholder ranges (offset, length) covering the entire - /// replacement including structural tokens. Used by vLLM (which filters via is_embed). - placeholders: Vec, - /// Patch-only placeholder ranges: contiguous runs of `im_token_id` within each - /// expansion. Used by sglang (which expects offsets aligned 1:1 with vision - /// encoder output). `None` when `im_token_id` is not set. - patch_offsets: Option>, + /// Full ranges covering each replacement, including structural tokens. + structural_ranges: Vec, + /// Contiguous `im_token_id` ranges aligned with encoder features. `None` + /// when the model does not declare a feature token ID. + feature_ranges: Option>, } /// Expand placeholder tokens in the token ID sequence. /// /// For each placeholder token found, replace it with the expanded token sequence /// from the corresponding `PromptReplacement`. Also track both the full structural -/// placeholder ranges and patch-only offsets (contiguous runs of `im_token_id`) +/// structural ranges and encoder-feature ranges (runs of `im_token_id`) /// in a single pass — no extra iteration needed. fn expand_tokens( token_ids: &[u32], @@ -877,45 +997,69 @@ fn expand_tokens( warn!("Could not resolve placeholder token ID; skipping token expansion"); return ExpandedTokens { token_ids: token_ids.to_vec(), - placeholders: vec![], - patch_offsets: None, + structural_ranges: vec![], + feature_ranges: None, }; }; - let mut expanded = Vec::with_capacity(token_ids.len()); - let mut placeholders = Vec::new(); - let mut patch_offsets: Option> = im_token_id.map(|_| Vec::new()); + let replacement_extra_capacity = replacements + .iter() + .try_fold(0usize, |acc, repl| { + acc.checked_add(repl.tokens.len().saturating_sub(1)) + }) + .unwrap_or(0); + let expanded_capacity = token_ids + .len() + .checked_add(replacement_extra_capacity) + .unwrap_or(token_ids.len()); + let mut expanded = Vec::with_capacity(expanded_capacity); + let mut structural_ranges = Vec::with_capacity(replacements.len()); + let mut feature_ranges: Option> = + im_token_id.map(|_| Vec::with_capacity(replacements.len())); let mut replacement_idx = 0; for &token in token_ids { if token == placeholder_id && replacement_idx < replacements.len() { let repl = &replacements[replacement_idx]; let offset = expanded.len(); - - // Track patch-only runs while extending - if let (Some(im_id), Some(ref mut offsets)) = (im_token_id, &mut patch_offsets) { - let mut run_start: Option = None; - for (i, &t) in repl.tokens.iter().enumerate() { - let pos = (offset + i) as u32; - if t as u32 == im_id { - if run_start.is_none() { - run_start = Some(pos); + let repl_len = repl.tokens.len(); + let mut repeated_patch_token: Option = None; + + // Track encoder-feature runs while extending. + if let (Some(im_id), Some(ref mut offsets)) = (im_token_id, &mut feature_ranges) { + if matches!(repl.tokens.first(), Some(&token) if token as u32 == im_id) + && repl.tokens.iter().all(|&token| token as u32 == im_id) + { + offsets.push((offset as u32, repl_len as u32)); + repeated_patch_token = Some(im_id); + } else { + let mut run_start: Option = None; + for (i, &t) in repl.tokens.iter().enumerate() { + let pos = (offset + i) as u32; + if t as u32 == im_id { + if run_start.is_none() { + run_start = Some(pos); + } + } else if let Some(s) = run_start { + offsets.push((s, pos - s)); + run_start = None; } - } else if let Some(s) = run_start { - offsets.push((s, pos - s)); - run_start = None; } - } - if let Some(s) = run_start { - offsets.push((s, (offset + repl.tokens.len()) as u32 - s)); + if let Some(s) = run_start { + offsets.push((s, (offset + repl_len) as u32 - s)); + } } } // PromptReplacement uses TokenId = i32, convert to u32 - expanded.extend(repl.tokens.iter().map(|&t| t as u32)); - placeholders.push(PlaceholderRange { + if let Some(token) = repeated_patch_token { + expanded.resize(expanded.len() + repl_len, token); + } else { + expanded.extend(repl.tokens.iter().map(|&t| t as u32)); + } + structural_ranges.push(PlaceholderRange { offset, - length: repl.tokens.len(), + length: repl_len, }); replacement_idx += 1; } else { @@ -933,969 +1077,131 @@ fn expand_tokens( ExpandedTokens { token_ids: expanded, - placeholders, - patch_offsets, - } -} - -// --------------------------------------------------------------------------- -// Assembly: convert MultimodalIntermediate → backend-specific MultimodalData -// --------------------------------------------------------------------------- - -/// Assemble backend-specific multimodal data from the intermediate. -/// -/// Called in request_building after worker selection, when the backend is known. -#[expect( - clippy::unreachable, - reason = "MLX multimodal rejected by caller before reaching here" -)] -pub(crate) fn assemble_multimodal_data( - intermediate: MultimodalIntermediate, - client: &GrpcClient, - workers: Option<&WorkerSelection>, -) -> Result { - match intermediate { - MultimodalIntermediate::Precomputed(precomputed) => match client { - GrpcClient::Sglang(_) => { - ensure_image_only(&precomputed, "SGLang")?; - Ok(MultimodalData::Sglang(assemble_sglang(precomputed))) - } - GrpcClient::Vllm(_) => { - ensure_image_only(&precomputed, "vLLM")?; - Ok(MultimodalData::Vllm(assemble_vllm(precomputed))) - } - GrpcClient::Trtllm(_) => { - ensure_image_only(&precomputed, "TRT-LLM")?; - Ok(MultimodalData::Trtllm(assemble_trtllm(precomputed))) - } - GrpcClient::TokenSpeed(_) => Ok(MultimodalData::TokenSpeed(assemble_tokenspeed( - precomputed, - workers, - )?)), - GrpcClient::Mlx(_) => unreachable!( - "caller rejects multimodal for MLX in build_chat_request/build_messages_request" - ), - }, - } -} - -fn ensure_image_only( - intermediate: &PrecomputedMultimodalIntermediate, - backend: &str, -) -> Result<()> { - if intermediate.modality != Modality::Image { - return Err(anyhow::anyhow!( - "{backend} multimodal path currently supports image inputs only; got {}", - intermediate.modality - )); - } - Ok(()) -} - -fn assemble_sglang(intermediate: PrecomputedMultimodalIntermediate) -> SglangMultimodalData { - let (pixel_values, pixel_values_shape) = serialize_encoder_input(&intermediate.preprocessed); - let model_specific_tensors = serialize_model_specific(intermediate.preprocessed.model_specific); - let image_data = intermediate - .images - .iter() - .map(|f| f.raw_bytes.to_vec()) - .collect(); - // Use patch-only offsets when available and non-empty; fall back to full structural ranges. - let mm_placeholders = intermediate - .patch_offsets - .filter(|offsets| !offsets.is_empty()) - .unwrap_or_else(|| { - intermediate - .placeholders - .iter() - .map(|p| (p.offset as u32, p.length as u32)) - .collect() - }); - - SglangMultimodalData { - image_data, - pixel_values, - pixel_values_shape, - model_specific_tensors, - im_token_id: intermediate.placeholder_token_id, - mm_placeholders, - } -} - -fn assemble_vllm(intermediate: PrecomputedMultimodalIntermediate) -> VllmMultimodalData { - let (pixel_values, pixel_values_shape) = serialize_encoder_input(&intermediate.preprocessed); - let model_specific_tensors = serialize_model_specific(intermediate.preprocessed.model_specific); - let mm_hashes = intermediate.images.iter().map(|f| f.hash.clone()).collect(); - let mm_placeholders = intermediate - .placeholders - .iter() - .map(|p| (p.offset as u32, p.length as u32)) - .collect(); - let batched_keys = PreprocessedEncoderInputs::batched_keys(&intermediate.field_layouts); - let flat_keys = PreprocessedEncoderInputs::flat_keys(&intermediate.field_layouts); - - VllmMultimodalData { - pixel_values, - pixel_values_shape, - model_specific_tensors, - im_token_id: intermediate.placeholder_token_id, - mm_placeholders, - mm_hashes, - batched_keys, - flat_keys, - keep_on_cpu_keys: intermediate.keep_on_cpu_keys, + structural_ranges, + feature_ranges, } } -fn assemble_trtllm(intermediate: PrecomputedMultimodalIntermediate) -> TrtllmMultimodalData { - let image_data = intermediate - .images - .iter() - .map(|f| f.raw_bytes.to_vec()) - .collect(); - TrtllmMultimodalData { image_data } -} +#[cfg(test)] +mod tests { + use std::{ + fs, + io::{Read, Seek, SeekFrom, Write}, + mem::size_of, + }; -fn assemble_tokenspeed( - intermediate: PrecomputedMultimodalIntermediate, - workers: Option<&WorkerSelection>, -) -> Result { - let log_timing = log_mm_timing_enabled(); - let total_started = Instant::now(); - // Resolve the multimodal tensor transport once per request: `shm` always on, - // `auto` only when the worker is verified to share /dev/shm (matching - // namespace token), otherwise inline. See `worker_shares_dev_shm`. - let shm_enabled = resolve_tokenspeed_shm_enabled(workers); - // Use patch-only offsets when available and non-empty; fall back to full structural ranges. - let encoder_input_dtype = tokenspeed_encoder_input_dtype(intermediate.modality, workers); - let patch_offsets = intermediate - .patch_offsets - .clone() - .filter(|offsets| !offsets.is_empty()) - .unwrap_or_default(); + use llm_multimodal::{DeferredNormalizedEncoderInput, ModelSpecificValue}; + use ndarray::{ArrayD, Axis, IxDyn, Slice}; + use openai_protocol::common::{ImageUrl, VideoUrl}; + use tempfile::TempDir; - let modality = match intermediate.modality { - Modality::Image => TokenSpeedModality::Image, - Modality::Video => TokenSpeedModality::Video, - Modality::Audio => TokenSpeedModality::Audio, - Modality::ImageEmbeds => TokenSpeedModality::Image, + use super::*; + use crate::routers::grpc::proto_wrapper::{ + cleanup_tokenspeed_shm_handles, tokenspeed_shm_dev_writable, write_tokenspeed_shm_with, + TokenSpeedModality, TokenSpeedTensorStorage, }; - let item_count = precomputed_multimodal_item_count(&intermediate)?; - // Build items imperatively so that if any step fails partway we can unlink - // the /dev/shm segments already created for prior items' encoder inputs - // (and this item's, once created). `?`/`collect` would drop those - // `TokenSpeedTensor::Shm` handles without ever reaching the send-path - // cleanup, leaking files until the next sweep. - let mut items: Vec = Vec::with_capacity(item_count); - for item_index in 0..item_count { - let item_encoder_input = match encoder_input_for_item( - &intermediate.preprocessed, - &intermediate.field_layouts, - item_index, - ) { - Ok(value) => value, - Err(error) => { - cleanup_tokenspeed_items_encoder_shm(&items, None); - return Err(error); - } - }; - let encoder_input_started = Instant::now(); - let encoder_input = serialize_array_as_tokenspeed_tensor( - &item_encoder_input, - &encoder_input_dtype, - shm_enabled, + #[test] + #[cfg(target_os = "linux")] + fn local_shm_namespace_id_resolves_on_linux() { + // /proc/.../boot_id and /dev/shm both exist on the Linux CI/runtime + // image, so the token must resolve to `:`. If it ever + // returned None, `auto` would silently never enable SHM. + let id = local_shm_namespace_id().expect("shm namespace id should resolve on Linux"); + assert!( + id.contains(':'), + "token must be :, got {id:?}" ); - let encoder_input_serialize_ms = encoder_input_started.elapsed().as_secs_f64() * 1000.0; - let model_specific_started = Instant::now(); - let model_specific_tensors = match serialize_model_specific_for_item( - &intermediate.preprocessed.model_specific, - &intermediate.field_layouts, - item_index, - ) { - Ok(value) => value, - Err(error) => { - // `encoder_input` (possibly SHM) was created for this item but the - // item isn't built; clean it plus all prior items. - cleanup_tokenspeed_items_encoder_shm(&items, Some(&encoder_input)); - return Err(error); - } - }; - let model_specific_serialize_ms = model_specific_started.elapsed().as_secs_f64() * 1000.0; - let mm_placeholders = - placeholders_for_item(item_index, &intermediate.placeholders, &patch_offsets); - let content_hash = content_hash_for_item(intermediate.modality, &intermediate, item_index); - - if log_timing { - info!( - modality = ?modality, - item_index, - encoder_input_dtype = %encoder_input.dtype, - encoder_input_bytes = encoder_input.nbytes(), - encoder_input_shape = ?encoder_input.shape, - model_specific_tensor_count = model_specific_tensors.len(), - encoder_input_serialize_ms, - model_specific_serialize_ms, - "smg_mm_timing assemble_tokenspeed_item" - ); - } - - items.push(TokenSpeedMultimodalItem { - modality, - encoder_input, - model_specific_tensors, - placeholder_token_id: intermediate.placeholder_token_id, - mm_placeholders, - content_hash, - }); - } - - if log_timing { - info!( - modality = ?modality, - item_count = items.len(), - total_ms = total_started.elapsed().as_secs_f64() * 1000.0, - "smg_mm_timing assemble_tokenspeed" + let dev = id.rsplit(':').next().unwrap(); + assert!( + dev.parse::().is_ok(), + "st_dev component must be numeric, got {id:?}" ); } - Ok(TokenSpeedMultimodalData { items, shm_enabled }) -} - -fn precomputed_multimodal_item_count( - intermediate: &PrecomputedMultimodalIntermediate, -) -> Result { - let media_count = match intermediate.modality { - Modality::Image | Modality::ImageEmbeds => intermediate.images.len(), - Modality::Video => intermediate.videos.len(), - Modality::Audio => 0, - }; - let token_count = intermediate.preprocessed.feature_token_counts.len(); - let placeholder_count = intermediate.placeholders.len(); - let item_count = token_count.max(media_count).max(placeholder_count); - anyhow::ensure!( - item_count > 0, - "precomputed multimodal assembly requires at least one item" - ); - if media_count > 0 { - anyhow::ensure!( - media_count == item_count, - "precomputed multimodal assembly media count mismatch: modality={}, media_count={media_count}, item_count={item_count}", - intermediate.modality + #[test] + fn tokenspeed_transport_default_uses_video_auto_only() { + assert_eq!( + effective_tokenspeed_transport_mode(Modality::Image, ""), + "inline" + ); + assert_eq!( + effective_tokenspeed_transport_mode(Modality::ImageEmbeds, ""), + "inline" + ); + assert_eq!( + effective_tokenspeed_transport_mode(Modality::Audio, ""), + "inline" + ); + assert_eq!( + effective_tokenspeed_transport_mode(Modality::Video, ""), + "auto" ); } - anyhow::ensure!( - token_count == item_count, - "precomputed multimodal assembly token count mismatch: modality={}, token_count={token_count}, item_count={item_count}", - intermediate.modality - ); - anyhow::ensure!( - placeholder_count == item_count, - "precomputed multimodal assembly placeholder count mismatch: modality={}, placeholder_count={placeholder_count}, item_count={item_count}", - intermediate.modality - ); - Ok(item_count) -} -fn encoder_input_for_item( - preprocessed: &PreprocessedEncoderInputs, - field_layouts: &HashMap, - item_index: usize, -) -> Result> { - // The field layout key remains "pixel_values" because it mirrors the - // HuggingFace/vLLM vision kwargs contract. Internally this tensor is the - // modality encoder input we pass to TokenSpeed. - let layout = field_layouts - .get("pixel_values") - .unwrap_or(&FieldLayout::Batched); - match layout { - FieldLayout::Batched => slice_array_axis0(&preprocessed.encoder_input, item_index, 1), - FieldLayout::Flat { sizes_key } => { - let sizes = tensor_sizes_from_model_specific(&preprocessed.model_specific, sizes_key)?; - let (start, len) = item_span(&sizes, item_index)?; - slice_array_axis0(&preprocessed.encoder_input, start, len) - } + #[test] + fn tokenspeed_transport_explicit_mode_overrides_modality_default() { + assert_eq!( + effective_tokenspeed_transport_mode(Modality::Image, "auto"), + "auto" + ); + assert_eq!( + effective_tokenspeed_transport_mode(Modality::Video, "inline"), + "inline" + ); + assert_eq!( + effective_tokenspeed_transport_mode(Modality::Video, "shm"), + "shm" + ); } -} -fn serialize_model_specific_for_item( - model_specific: &HashMap, - field_layouts: &HashMap, - item_index: usize, -) -> Result> { - let mut serialized = HashMap::with_capacity(model_specific.len()); - for (key, value) in model_specific { - let item_value = match field_layouts.get(key) { - Some(FieldLayout::Batched) => value - .slice_first_dim(item_index, 1) - .with_context(|| format!("failed to slice model_specific tensor {key}"))?, - Some(FieldLayout::Flat { sizes_key }) => { - let sizes = tensor_sizes_from_model_specific(model_specific, sizes_key)?; - let (start, len) = item_span(&sizes, item_index)?; - value - .slice_first_dim(start, len) - .with_context(|| format!("failed to slice flat model_specific tensor {key}"))? - } - None => value.clone(), + #[test] + fn prepared_media_from_tracker_builds_one_exhaustive_variant() { + let image = Arc::new(ImageFrame::new( + image::DynamicImage::new_rgb8(1, 1), + bytes::Bytes::from_static(b"image"), + ImageDetail::Auto, + llm_multimodal::ImageSource::InlineBytes, + "image-hash".to_string(), + )); + let output = TrackerOutput { + data: HashMap::from([(Modality::Image, vec![TrackedMedia::Image(image.clone())])]), + uuids: HashMap::new(), }; - if let Some(tensor) = model_specific_to_tensor_bytes(&item_value) { - serialized.insert(key.clone(), tensor); - } else { - warn!(tensor_key = %key, "Dropping unsupported model_specific value during multimodal serialization"); - } - } - Ok(serialized) -} - -fn placeholders_for_item( - item_index: usize, - placeholders: &[PlaceholderRange], - patch_offsets: &[(u32, u32)], -) -> Vec<(u32, u32)> { - let Some(placeholder) = placeholders.get(item_index) else { - return Vec::new(); - }; - let start = placeholder.offset as u32; - let end = start + placeholder.length as u32; - let item_patch_offsets = patch_offsets - .iter() - .copied() - .filter(|(offset, length)| *offset >= start && offset.saturating_add(*length) <= end) - .collect::>(); - if item_patch_offsets.is_empty() { - vec![(start, end - start)] - } else { - item_patch_offsets - } -} -fn content_hash_for_item( - modality: Modality, - intermediate: &PrecomputedMultimodalIntermediate, - item_index: usize, -) -> Vec { - match modality { - Modality::Image | Modality::ImageEmbeds => intermediate - .images - .get(item_index) - .map(|image| hash_hex_strings(std::iter::once(image.hash.as_str()))) - .unwrap_or_default(), - Modality::Video => intermediate - .videos - .get(item_index) - .map(|video| hash_hex_strings(std::iter::once(video.hash.as_str()))) - .unwrap_or_default(), - Modality::Audio => Vec::new(), + let prepared = prepared_media_from_tracker(output).unwrap(); + let PreparedMedia::Images(images) = prepared else { + panic!("expected prepared image media"); + }; + assert_eq!(images.len(), 1); + assert!(Arc::ptr_eq(&images[0], &image)); } -} - -fn slice_array_axis0(array: &ArrayD, start: usize, len: usize) -> Result> { - let end = start - .checked_add(len) - .ok_or_else(|| anyhow::anyhow!("array slice range overflow"))?; - let rows = array.shape().first().copied().unwrap_or(0); - anyhow::ensure!( - end <= rows, - "array first-dimension slice {start}..{end} exceeds {rows}" - ); - Ok(array - .slice_axis(Axis(0), Slice::from(start..end)) - .to_owned()) -} - -fn tensor_sizes_from_model_specific( - model_specific: &HashMap, - key: &str, -) -> Result> { - let value = model_specific - .get(key) - .ok_or_else(|| anyhow::anyhow!("missing flat sizes tensor {key}"))?; - value - .as_flat_sizes() - .with_context(|| format!("invalid flat sizes tensor {key}")) -} -fn item_span(sizes: &[usize], item_index: usize) -> Result<(usize, usize)> { - let len = *sizes - .get(item_index) - .ok_or_else(|| anyhow::anyhow!("missing flat size for item {item_index}"))?; - let start = sizes[..item_index] - .iter() - .try_fold(0usize, |acc, &size| acc.checked_add(size)) - .ok_or_else(|| anyhow::anyhow!("flat size offset overflow"))?; - Ok((start, len)) -} + #[test] + fn prepared_media_from_tracker_rejects_mixed_modalities() { + let image = Arc::new(ImageFrame::new( + image::DynamicImage::new_rgb8(1, 1), + bytes::Bytes::from_static(b"image"), + ImageDetail::Auto, + llm_multimodal::ImageSource::InlineBytes, + "image-hash".to_string(), + )); + let video = Arc::new(VideoClip::new( + vec![image::DynamicImage::new_rgb8(1, 1)], + bytes::Bytes::from_static(b"video"), + llm_multimodal::VideoSource::InlineBytes, + "video-hash".to_string(), + )); + let output = TrackerOutput { + data: HashMap::from([ + (Modality::Image, vec![TrackedMedia::Image(image)]), + (Modality::Video, vec![TrackedMedia::Video(video)]), + ]), + uuids: HashMap::new(), + }; -fn hash_hex_strings<'a>(hashes: impl Iterator) -> Vec { - let mut hasher = blake3::Hasher::new(); - for hash in hashes { - hasher.update(hash.as_bytes()); - } - hasher.finalize().as_bytes().to_vec() -} - -// --------------------------------------------------------------------------- -// Serialization helpers -// --------------------------------------------------------------------------- - -/// Serialize the primary encoder input ndarray to raw little-endian f32 bytes + shape. -fn serialize_encoder_input(preprocessed: &PreprocessedEncoderInputs) -> (Vec, Vec) { - serialize_array(&preprocessed.encoder_input) -} - -fn serialize_array(encoder_input: &ArrayD) -> (Vec, Vec) { - let encoder_bytes: Vec = if let Some(encoder_slice) = encoder_input - // Fast path only for C-contiguous arrays, whose memory order equals - // logical (row-major) order. A non-C-contiguous array (e.g. a - // Fortran-contiguous view) falls through to logical `.iter()` below; - // `as_slice_memory_order()` is deliberately NOT used as a fallback - // because it would serialize such arrays in the wrong dimension order. - .as_slice() - { - // Zero-copy reinterpret: &[f32] → &[u8] on little-endian (x86). - // This replaces the per-element flat_map(to_le_bytes) which was the - // #1 CPU hotspot (13% of SMG CPU in profiling). - #[cfg(target_endian = "little")] - { - let byte_slice: &[u8] = bytemuck::cast_slice(encoder_slice); - byte_slice.to_vec() - } - #[cfg(not(target_endian = "little"))] - { - encoder_slice.iter().flat_map(|v| v.to_le_bytes()).collect() - } - } else { - // Non-C-contiguous array: `.iter()` walks in logical (row-major) order, - // which matches the shape. - encoder_input.iter().flat_map(|v| v.to_le_bytes()).collect() - }; - (encoder_bytes, array_shape(encoder_input)) -} - -/// Serialize encoder input to the requested wire dtype. -fn serialize_array_as_tokenspeed_tensor( - encoder_input: &ArrayD, - dtype: &str, - shm_enabled: bool, -) -> TokenSpeedTensor { - let dtype = match canonical_float_dtype(dtype).as_deref() { - Some("float32") => "float32".to_string(), - Some("bfloat16") => "bfloat16".to_string(), - Some("float16") => "float16".to_string(), - _ => { - warn!( - dtype, - "Unsupported TokenSpeed encoder input dtype; falling back to float32" - ); - "float32".to_string() - } - }; - let shape = array_shape(encoder_input); - let element_size = if dtype == "bfloat16" || dtype == "float16" { - size_of::() - } else { - size_of::() - }; - let nbytes = encoder_input.len() * element_size; - - if shm_enabled && nbytes >= tokenspeed_mm_shm_min_bytes() { - let started = Instant::now(); - match write_tokenspeed_shm_with(nbytes, |file| { - write_array_as_dtype(file, encoder_input, &dtype) - }) { - Ok(handle) => { - if log_mm_timing_enabled() { - info!( - nbytes, - elapsed_ms = started.elapsed().as_secs_f64() * 1000.0, - "smg_mm_timing tokenspeed_shm_write_direct" - ); - } - return TokenSpeedTensor::shm(handle, shape, dtype); - } - Err(error) => { - use crate::observability::metrics::Metrics; - warn!( - ?error, - nbytes, - dtype = %dtype, - "Failed to write TokenSpeed encoder input directly to SHM; falling back to bytes path" - ); - Metrics::record_mm_shm_write_failure("tokenspeed"); - } - } - } - - let (data, shape, dtype) = serialize_array_as_dtype(encoder_input, &dtype); - TokenSpeedTensor::inline(data, shape, dtype) -} - -fn write_array_as_dtype( - writer: &mut impl Write, - encoder_input: &ArrayD, - dtype: &str, -) -> std::io::Result<()> { - match dtype { - "float32" => write_array_as_f32(writer, encoder_input), - "bfloat16" => write_array_as_u16(writer, encoder_input, f32_to_bf16_bits), - "float16" => write_array_as_u16(writer, encoder_input, f32_to_f16_bits), - other => Err(std::io::Error::new( - std::io::ErrorKind::InvalidInput, - format!("unsupported TokenSpeed encoder input dtype: {other}"), - )), - } -} - -fn write_array_as_f32(writer: &mut impl Write, encoder_input: &ArrayD) -> std::io::Result<()> { - if let Some(encoder_slice) = encoder_input - // Fast path only for C-contiguous arrays, whose memory order equals - // logical (row-major) order. A non-C-contiguous array (e.g. a - // Fortran-contiguous view) falls through to logical `.iter()` below; - // `as_slice_memory_order()` is deliberately NOT used as a fallback - // because it would serialize such arrays in the wrong dimension order. - .as_slice() - { - return write_f32_slice(writer, encoder_slice); - } - - for value in encoder_input { - writer.write_all(&value.to_le_bytes())?; - } - Ok(()) -} - -fn write_f32_slice(writer: &mut impl Write, values: &[f32]) -> std::io::Result<()> { - #[cfg(target_endian = "little")] - { - writer.write_all(bytemuck::cast_slice(values)) - } - #[cfg(not(target_endian = "little"))] - { - for value in values { - writer.write_all(&value.to_le_bytes())?; - } - Ok(()) - } -} - -fn write_array_as_u16( - writer: &mut impl Write, - encoder_input: &ArrayD, - convert: F, -) -> std::io::Result<()> -where - F: Fn(f32) -> u16 + Copy, -{ - // Convert in bounded chunks so peak memory stays at ~CHUNK_VALUES u16s - // regardless of tensor size, on both the contiguous and strided paths. - const CHUNK_VALUES: usize = 256 * 1024; - - if let Some(encoder_slice) = encoder_input - // Fast path only for C-contiguous arrays, whose memory order equals - // logical (row-major) order. A non-C-contiguous array (e.g. a - // Fortran-contiguous view) falls through to logical `.iter()` below; - // `as_slice_memory_order()` is deliberately NOT used as a fallback - // because it would serialize such arrays in the wrong dimension order. - .as_slice() - { - let mut converted: Vec = Vec::with_capacity(CHUNK_VALUES.min(encoder_slice.len())); - for chunk in encoder_slice.chunks(CHUNK_VALUES) { - converted.clear(); - converted.extend(chunk.iter().map(|&value| convert(value))); - #[cfg(target_endian = "little")] - { - writer.write_all(bytemuck::cast_slice(converted.as_slice()))?; - } - #[cfg(not(target_endian = "little"))] - { - for value in &converted { - writer.write_all(&value.to_le_bytes())?; - } - } - } - return Ok(()); - } - - let mut converted = Vec::with_capacity(CHUNK_VALUES); - let mut flush = |converted: &mut Vec| -> std::io::Result<()> { - if converted.is_empty() { - return Ok(()); - } - #[cfg(target_endian = "little")] - { - writer.write_all(bytemuck::cast_slice(converted.as_slice()))?; - } - #[cfg(not(target_endian = "little"))] - { - for value in converted.iter() { - writer.write_all(&value.to_le_bytes())?; - } - } - converted.clear(); - Ok(()) - }; - - for &value in encoder_input { - converted.push(convert(value)); - if converted.len() == CHUNK_VALUES { - flush(&mut converted)?; - } - } - flush(&mut converted) -} - -fn serialize_array_as_dtype( - encoder_input: &ArrayD, - dtype: &str, -) -> (Vec, Vec, String) { - match canonical_float_dtype(dtype).as_deref() { - Some("float32") => { - let (data, shape) = serialize_array(encoder_input); - (data, shape, "float32".to_string()) - } - Some("bfloat16") => ( - serialize_array_as_u16_bytes(encoder_input, f32_to_bf16_bits), - array_shape(encoder_input), - "bfloat16".to_string(), - ), - Some("float16") => ( - serialize_array_as_u16_bytes(encoder_input, f32_to_f16_bits), - array_shape(encoder_input), - "float16".to_string(), - ), - _ => { - warn!( - dtype, - "Unsupported TokenSpeed encoder input dtype; falling back to float32" - ); - let (data, shape) = serialize_array(encoder_input); - (data, shape, "float32".to_string()) - } - } -} - -fn serialize_array_as_u16_bytes(encoder_input: &ArrayD, convert: F) -> Vec -where - F: Fn(f32) -> u16 + Copy, -{ - let element_count = encoder_input.len(); - let mut converted = Vec::with_capacity(element_count); - - if let Some(encoder_slice) = encoder_input - // Fast path only for C-contiguous arrays, whose memory order equals - // logical (row-major) order. A non-C-contiguous array (e.g. a - // Fortran-contiguous view) falls through to logical `.iter()` below; - // `as_slice_memory_order()` is deliberately NOT used as a fallback - // because it would serialize such arrays in the wrong dimension order. - .as_slice() - { - converted.extend(encoder_slice.iter().map(|&value| convert(value))); - } else { - converted.extend(encoder_input.iter().map(|&value| convert(value))); - } - - #[cfg(target_endian = "little")] - { - bytemuck::cast_slice(&converted).to_vec() - } - #[cfg(not(target_endian = "little"))] - { - let mut bytes = Vec::with_capacity(element_count * std::mem::size_of::()); - for value in converted { - bytes.extend_from_slice(&value.to_le_bytes()); - } - bytes - } -} - -fn tokenspeed_encoder_input_dtype(modality: Modality, workers: Option<&WorkerSelection>) -> String { - if let Some(dtype) = tokenspeed_encoder_input_dtype_from_env(modality) { - return dtype; - } - if let Some(dtype) = tokenspeed_encoder_input_dtype_from_worker(workers) { - return dtype; - } - "float32".to_string() -} - -fn tokenspeed_encoder_input_dtype_from_env(modality: Modality) -> Option { - static IMAGE_DTYPE: OnceLock> = OnceLock::new(); - static VIDEO_DTYPE: OnceLock> = OnceLock::new(); - static AUDIO_DTYPE: OnceLock> = OnceLock::new(); - static DEFAULT_DTYPE: OnceLock> = OnceLock::new(); - - let modality_dtype = match modality { - Modality::Image | Modality::ImageEmbeds => { - cached_env_dtype(&IMAGE_DTYPE, "SMG_TOKENSPEED_IMAGE_ENCODER_INPUT_DTYPE") - } - Modality::Video => { - cached_env_dtype(&VIDEO_DTYPE, "SMG_TOKENSPEED_VIDEO_ENCODER_INPUT_DTYPE") - } - Modality::Audio => { - cached_env_dtype(&AUDIO_DTYPE, "SMG_TOKENSPEED_AUDIO_ENCODER_INPUT_DTYPE") - } - }; - modality_dtype - .or_else(|| cached_env_dtype(&DEFAULT_DTYPE, "SMG_TOKENSPEED_ENCODER_INPUT_DTYPE")) -} - -fn cached_env_dtype(cell: &'static OnceLock>, name: &str) -> Option { - cell.get_or_init(|| std::env::var(name).ok().filter(|dtype| !dtype.is_empty())) - .clone() -} - -fn tokenspeed_encoder_input_dtype_from_worker(workers: Option<&WorkerSelection>) -> Option { - let worker = match workers? { - WorkerSelection::Single { worker } => worker, - WorkerSelection::Dual { prefill, .. } => prefill, - }; - worker - .metadata() - .spec - .labels - .get("multimodal_encoder_dtype") - .filter(|dtype| !dtype.is_empty()) - .cloned() -} - -/// Resolve whether large multimodal tensors should use the SHM transport for -/// this request. `shm` = always (legacy explicit opt-in); `auto` = only when the -/// worker is known to share SMG's `/dev/shm`; anything else (including unset or -/// `inline`) keeps the inline gRPC path. -fn resolve_tokenspeed_shm_enabled(workers: Option<&WorkerSelection>) -> bool { - let mode = tokenspeed_mm_tensor_transport_mode(); - log_tokenspeed_transport_config_once(&mode); - match mode.as_str() { - // SHM only ever happens when SMG can actually write /dev/shm. - "shm" => tokenspeed_shm_dev_writable(), - "auto" => worker_shares_dev_shm(workers) && tokenspeed_shm_dev_writable(), - "" | "inline" => false, - other => { - log_unknown_tokenspeed_transport_once(other); - false - } - } -} - -fn log_tokenspeed_transport_config_once(mode: &str) { - static LOGGED: OnceLock<()> = OnceLock::new(); - LOGGED.get_or_init(|| { - info!( - mode, - shm_min_bytes = tokenspeed_mm_shm_min_bytes(), - dev_writable = tokenspeed_shm_dev_writable(), - "TokenSpeed multimodal tensor transport configured" - ); - }); -} - -fn log_unknown_tokenspeed_transport_once(value: &str) { - static WARNED: OnceLock<()> = OnceLock::new(); - WARNED.get_or_init(|| { - warn!( - value, - "Unknown SMG_TOKENSPEED_MM_TENSOR_TRANSPORT value; expected inline|shm|auto, using inline" - ); - }); -} - -/// Whether the worker is *verified* to share SMG's `/dev/shm`, making the SHM -/// transport safe under `auto`. -/// -/// Rather than inferring locality from the worker URL (TCP loopback proves only -/// network locality, not a shared `/dev/shm`), the worker advertises its -/// `/dev/shm` filesystem identity (`:`) via -/// `GetServerInfo`, which discovery stores in the worker's `shm_namespace_id` -/// label. Two processes share `/dev/shm` iff these tokens match: `boot_id` pins -/// the host, and `st_dev` is the tmpfs superblock device, identical whenever the -/// same tmpfs backs both `/dev/shm` mounts — including separate containers that -/// share it via `--ipc`/bind-mount (where mount-namespace inodes differ but the -/// underlying superblock is the same). We compare the worker's token to ours: -/// equal ⇒ shared. A missing/empty token or any mismatch is treated as -/// non-sharing, so `auto` safely falls back to inline. -fn worker_shares_dev_shm(workers: Option<&WorkerSelection>) -> bool { - let Some(local) = local_shm_namespace_id() else { - return false; - }; - let worker = match workers { - Some(WorkerSelection::Single { worker }) => worker, - Some(WorkerSelection::Dual { prefill, .. }) => prefill, - None => return false, - }; - worker - .metadata() - .spec - .labels - .get("shm_namespace_id") - .is_some_and(|id| !id.is_empty() && id == local) -} - -/// This process's `/dev/shm` filesystem identity: `:`. -/// `boot_id` pins the host (it is not namespaced) and `st_dev` is the tmpfs -/// superblock device backing `/dev/shm`; together they identify the tmpfs so two -/// processes sharing it (even across containers via `--ipc`/bind-mount) produce -/// the same token. Computed once; `None` if it can't be determined (then `auto` -/// stays inline). -fn local_shm_namespace_id() -> Option<&'static str> { - static ID: OnceLock> = OnceLock::new(); - ID.get_or_init(compute_shm_namespace_id).as_deref() -} - -#[cfg(unix)] -fn compute_shm_namespace_id() -> Option { - use std::os::unix::fs::MetadataExt; - let boot_id = std::fs::read_to_string("/proc/sys/kernel/random/boot_id").ok()?; - let shm_dev = std::fs::metadata("/dev/shm").ok()?.dev(); - Some(format!("{}:{shm_dev}", boot_id.trim())) -} - -#[cfg(not(unix))] -fn compute_shm_namespace_id() -> Option { - None -} - -fn canonical_float_dtype(dtype: &str) -> Option { - match dtype.trim().to_ascii_lowercase().as_str() { - "float32" | "fp32" | "f32" => Some("float32".to_string()), - "bfloat16" | "bf16" => Some("bfloat16".to_string()), - "float16" | "fp16" | "f16" | "half" => Some("float16".to_string()), - _ => None, - } -} - -fn array_shape(encoder_input: &ArrayD) -> Vec { - encoder_input.shape().iter().map(|&d| d as u32).collect() -} - -#[inline] -fn f32_to_bf16_bits(value: f32) -> u16 { - let bits = value.to_bits(); - let lsb = (bits >> 16) & 1; - let rounding_bias = 0x7fff + lsb; - (bits.wrapping_add(rounding_bias) >> 16) as u16 -} - -#[inline] -fn f32_to_f16_bits(value: f32) -> u16 { - let bits = value.to_bits(); - let sign = ((bits >> 16) & 0x8000) as u16; - let exp = ((bits >> 23) & 0xff) as i32; - let mant = bits & 0x7fffff; - - if exp == 0xff { - return if mant == 0 { - sign | 0x7c00 - } else { - sign | 0x7e00 - }; - } - - let half_exp = exp - 127 + 15; - if half_exp >= 0x1f { - return sign | 0x7c00; - } - if half_exp <= 0 { - if half_exp < -10 { - return sign; - } - let mantissa = mant | 0x800000; - let shift = (14 - half_exp) as u32; - let mut half_mant = (mantissa >> shift) as u16; - let round_bit = (mantissa >> (shift - 1)) & 1; - let sticky = mantissa & ((1u32 << (shift - 1)) - 1); - if round_bit != 0 && (sticky != 0 || (half_mant & 1) != 0) { - half_mant += 1; - } - return sign | half_mant; - } - - let mut half = sign | ((half_exp as u16) << 10) | ((mant >> 13) as u16); - let round = mant & 0x1fff; - if round > 0x1000 || (round == 0x1000 && (half & 1) != 0) { - half += 1; - } - half -} - -/// Serialize model-specific values to TensorBytes, consuming the map to avoid key clones. -fn serialize_model_specific( - model_specific: HashMap, -) -> HashMap { - model_specific - .into_iter() - .filter_map(|(key, value)| match model_specific_to_tensor_bytes(&value) { - Some(tensor) => Some((key, tensor)), - None => { - warn!(tensor_key = %key, "Dropping unsupported model_specific value during multimodal serialization"); - None - } - }) - .collect() -} - -/// Convert a model-specific value to backend-agnostic TensorBytes. -fn model_specific_to_tensor_bytes(value: &ModelSpecificValue) -> Option { - match value { - ModelSpecificValue::Tensor { data, shape } => Some(TensorBytes { - data: data.iter().flat_map(|v| v.to_le_bytes()).collect(), - shape: shape.iter().map(|&d| d as u32).collect(), - dtype: "float32".to_string(), - }), - ModelSpecificValue::IntTensor { data, shape } => Some(TensorBytes { - data: data.iter().flat_map(|v| v.to_le_bytes()).collect(), - shape: shape.iter().map(|&d| d as u32).collect(), - dtype: "int64".to_string(), - }), - ModelSpecificValue::UintTensor { data, shape } => Some(TensorBytes { - data: data.iter().flat_map(|v| v.to_le_bytes()).collect(), - shape: shape.iter().map(|&d| d as u32).collect(), - dtype: "uint32".to_string(), - }), - ModelSpecificValue::UintVec(v) => Some(TensorBytes { - data: v.iter().flat_map(|val| val.to_le_bytes()).collect(), - shape: vec![v.len() as u32], - dtype: "uint32".to_string(), - }), - ModelSpecificValue::IntVec(v) => Some(TensorBytes { - data: v.iter().flat_map(|val| val.to_le_bytes()).collect(), - shape: vec![v.len() as u32], - dtype: "int64".to_string(), - }), - ModelSpecificValue::FloatVec(v) => Some(TensorBytes { - data: v.iter().flat_map(|val| val.to_le_bytes()).collect(), - shape: vec![v.len() as u32], - dtype: "float32".to_string(), - }), - _ => None, - } -} - -#[cfg(test)] -mod tests { - use std::{fs, mem::size_of}; - - use ndarray::IxDyn; - use openai_protocol::common::{ImageUrl, VideoUrl}; - use tempfile::TempDir; - - use super::*; - - #[test] - #[cfg(target_os = "linux")] - fn local_shm_namespace_id_resolves_on_linux() { - // /proc/.../boot_id and /dev/shm both exist on the Linux CI/runtime - // image, so the token must resolve to `:`. If it ever - // returned None, `auto` would silently never enable SHM. - let id = local_shm_namespace_id().expect("shm namespace id should resolve on Linux"); - assert!( - id.contains(':'), - "token must be :, got {id:?}" - ); - let dev = id.rsplit(':').next().unwrap(); - assert!( - dev.parse::().is_ok(), - "st_dev component must be numeric, got {id:?}" - ); + let error = prepared_media_from_tracker(output).unwrap_err(); + assert!(error.to_string().contains("Mixed multimodal requests")); } #[test] @@ -2028,10 +1334,10 @@ mod tests { let result = expand_tokens(&token_ids, Some(100), None, &replacements); assert_eq!(result.token_ids, vec![1, 2, 50, 50, 50, 50, 3, 4]); - assert_eq!(result.placeholders.len(), 1); - assert_eq!(result.placeholders[0].offset, 2); - assert_eq!(result.placeholders[0].length, 4); - assert!(result.patch_offsets.is_none()); + assert_eq!(result.structural_ranges.len(), 1); + assert_eq!(result.structural_ranges[0].offset, 2); + assert_eq!(result.structural_ranges[0].length, 4); + assert!(result.feature_ranges.is_none()); } #[test] @@ -2040,8 +1346,8 @@ mod tests { let result = expand_tokens(&token_ids, None, None, &[]); assert_eq!(result.token_ids, vec![1, 2, 3]); - assert!(result.placeholders.is_empty()); - assert!(result.patch_offsets.is_none()); + assert!(result.structural_ranges.is_empty()); + assert!(result.feature_ranges.is_none()); } #[test] @@ -2063,11 +1369,11 @@ mod tests { let result = expand_tokens(&token_ids, Some(100), None, &replacements); assert_eq!(result.token_ids, vec![1, 50, 50, 2, 60, 60, 60, 3]); - assert_eq!(result.placeholders.len(), 2); - assert_eq!(result.placeholders[0].offset, 1); - assert_eq!(result.placeholders[0].length, 2); - assert_eq!(result.placeholders[1].offset, 4); - assert_eq!(result.placeholders[1].length, 3); + assert_eq!(result.structural_ranges.len(), 2); + assert_eq!(result.structural_ranges[0].offset, 1); + assert_eq!(result.structural_ranges[0].length, 2); + assert_eq!(result.structural_ranges[1].offset, 4); + assert_eq!(result.structural_ranges[1].length, 3); } #[test] @@ -2084,17 +1390,193 @@ mod tests { let result = expand_tokens(&token_ids, Some(100), Some(92), &replacements); // Full structural range - assert_eq!(result.placeholders.len(), 1); - assert_eq!(result.placeholders[0].offset, 1); - assert_eq!(result.placeholders[0].length, 9); + assert_eq!(result.structural_ranges.len(), 1); + assert_eq!(result.structural_ranges[0].offset, 1); + assert_eq!(result.structural_ranges[0].length, 9); // Patch-only offsets: two runs of token 92 - let patch = result.patch_offsets.unwrap(); + let patch = result.feature_ranges.unwrap(); assert_eq!(patch.len(), 2); assert_eq!(patch[0], (2, 3)); // offset=2, length=3 assert_eq!(patch[1], (6, 3)); // offset=6, length=3 } + #[test] + fn test_expand_tokens_patch_offsets_all_repeated_patch_tokens() { + let token_ids = vec![1, 100, 2]; + let replacements = vec![PromptReplacement { + modality: Modality::Image, + placeholder_token: "".to_string(), + tokens: vec![92, 92, 92, 92], + }]; + + let result = expand_tokens(&token_ids, Some(100), Some(92), &replacements); + + assert_eq!(result.token_ids, vec![1, 92, 92, 92, 92, 2]); + assert_eq!(result.structural_ranges[0].offset, 1); + assert_eq!(result.structural_ranges[0].length, 4); + assert_eq!(result.feature_ranges.unwrap(), vec![(1, 4)]); + } + + #[test] + fn test_placeholders_for_items_uses_patch_offsets_in_one_pass() { + let placeholders = vec![ + PlaceholderRange { + offset: 1, + length: 9, + }, + PlaceholderRange { + offset: 11, + length: 4, + }, + ]; + let patch_offsets = vec![(2, 3), (6, 3), (12, 2)]; + + let by_item = placeholders_for_items(&placeholders, &patch_offsets); + + assert_eq!(by_item, vec![vec![(2, 3), (6, 3)], vec![(12, 2)]]); + } + + #[test] + fn test_placeholders_for_items_single_item_fast_path_matches_patch_offsets() { + let placeholders = vec![PlaceholderRange { + offset: 1, + length: 9, + }]; + let patch_offsets = vec![(2, 3), (6, 3)]; + + let by_item = placeholders_for_items(&placeholders, &patch_offsets); + + assert_eq!(by_item, vec![vec![(2, 3), (6, 3)]]); + } + + #[test] + fn test_placeholders_for_items_single_patch_run_fast_path() { + let placeholders = vec![PlaceholderRange { + offset: 1, + length: 9, + }]; + let patch_offsets = vec![(2, 3)]; + + let by_item = placeholders_for_items(&placeholders, &patch_offsets); + + assert_eq!(by_item, vec![vec![(2, 3)]]); + } + + #[test] + fn test_placeholders_for_items_one_patch_run_per_item_fast_path() { + let placeholders = vec![ + PlaceholderRange { + offset: 1, + length: 4, + }, + PlaceholderRange { + offset: 8, + length: 6, + }, + ]; + let patch_offsets = vec![(1, 4), (9, 3)]; + + let by_item = placeholders_for_items(&placeholders, &patch_offsets); + + assert_eq!(by_item, vec![vec![(1, 4)], vec![(9, 3)]]); + } + + #[test] + fn test_placeholders_for_items_falls_back_to_full_ranges() { + let placeholders = vec![ + PlaceholderRange { + offset: 1, + length: 2, + }, + PlaceholderRange { + offset: 4, + length: 3, + }, + ]; + + let by_item = placeholders_for_items(&placeholders, &[]); + + assert_eq!(by_item, vec![vec![(1, 2)], vec![(4, 3)]]); + } + + #[test] + fn test_flat_item_spans_precomputes_prefix_offsets() { + let model_specific = HashMap::from([( + "patches_per_image".to_string(), + ModelSpecificValue::UintTensor { + data: vec![2, 3, 1], + shape: vec![3], + }, + )]); + let field_layouts = HashMap::from([( + "pixel_values".to_string(), + FieldLayout::flat("patches_per_image"), + )]); + + let spans = flat_item_spans(&model_specific, &field_layouts, 3).unwrap(); + + assert_eq!(spans["patches_per_image"], vec![(0, 2), (2, 3), (5, 1)]); + assert_eq!( + flat_item_span(&spans, "patches_per_image", 1).unwrap(), + (2, 3) + ); + } + + #[test] + fn test_validate_tokenspeed_item_spans_rejects_flat_under_consumption() { + let preprocessed = PreprocessedEncoderInputs { + encoder_input: EncoderInput::Dense( + ArrayD::from_shape_vec(IxDyn(&[4, 2]), vec![0.0; 8]).unwrap(), + ), + feature_token_counts: vec![1, 1], + item_sizes: vec![(1, 1), (1, 1)], + model_specific: HashMap::from([( + "patches_per_image".to_string(), + ModelSpecificValue::UintTensor { + data: vec![2, 1], + shape: vec![2], + }, + )]), + }; + let field_layouts = HashMap::from([ + ( + "pixel_values".to_string(), + FieldLayout::flat("patches_per_image"), + ), + ("patches_per_image".to_string(), FieldLayout::Batched), + ]); + let flat_spans = flat_item_spans(&preprocessed.model_specific, &field_layouts, 2).unwrap(); + + let error = validate_tokenspeed_item_spans(&preprocessed, &field_layouts, &flat_spans, 2) + .unwrap_err(); + + assert!(error + .to_string() + .contains("flat tensor pixel_values first dimension mismatch")); + } + + #[test] + fn test_validate_tokenspeed_item_spans_rejects_batched_extra_rows() { + let preprocessed = PreprocessedEncoderInputs { + encoder_input: EncoderInput::Dense( + ArrayD::from_shape_vec(IxDyn(&[3, 2]), vec![0.0; 6]).unwrap(), + ), + feature_token_counts: vec![1, 1], + item_sizes: vec![(1, 1), (1, 1)], + model_specific: HashMap::new(), + }; + let field_layouts = HashMap::new(); + let flat_spans = HashMap::new(); + + let error = validate_tokenspeed_item_spans(&preprocessed, &field_layouts, &flat_spans, 2) + .unwrap_err(); + + assert!(error + .to_string() + .contains("batched tensor pixel_values first dimension mismatch")); + } + #[test] fn test_parse_detail() { assert_eq!(parse_detail("auto"), Some(ImageDetail::Auto)); @@ -2123,11 +1605,13 @@ mod tests { ); let preprocessed = PreprocessedEncoderInputs { - encoder_input: ArrayD::from_shape_vec( - IxDyn(&[4, 2]), - vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], - ) - .unwrap(), + encoder_input: EncoderInput::Dense( + ArrayD::from_shape_vec( + IxDyn(&[4, 2]), + vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + ) + .unwrap(), + ), feature_token_counts: vec![2, 2], item_sizes: vec![(1, 1), (1, 1)], model_specific, @@ -2150,12 +1634,10 @@ mod tests { )), ]; - let intermediate = PrecomputedMultimodalIntermediate { - modality: Modality::Image, - preprocessed, - images, - videos: vec![], - placeholders: vec![ + let intermediate = MultimodalIntermediate { + media: PreparedMedia::Images(images), + preprocessed: Arc::new(preprocessed), + structural_ranges: vec![ PlaceholderRange { offset: 10, length: 2, @@ -2165,7 +1647,7 @@ mod tests { length: 2, }, ], - patch_offsets: Some(vec![(10, 2), (20, 2)]), + feature_ranges: Some(vec![(10, 2), (20, 2)]), placeholder_token_id: Some(151655), field_layouts: HashMap::from([ ( @@ -2175,7 +1657,7 @@ mod tests { ("patches_per_image".to_string(), FieldLayout::Batched), ("image_grid_thw".to_string(), FieldLayout::Batched), ]), - keep_on_cpu_keys: vec![], + cpu_resident_tensor_keys: vec![], }; let assembled = assemble_tokenspeed(intermediate, None).unwrap(); @@ -2231,11 +1713,13 @@ mod tests { ); let preprocessed = PreprocessedEncoderInputs { - encoder_input: ArrayD::from_shape_vec( - IxDyn(&[4, 2]), - vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], - ) - .unwrap(), + encoder_input: EncoderInput::Dense( + ArrayD::from_shape_vec( + IxDyn(&[4, 2]), + vec![1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0], + ) + .unwrap(), + ), feature_token_counts: vec![2, 2], item_sizes: vec![(1, 1), (1, 1)], model_specific, @@ -2256,12 +1740,10 @@ mod tests { )), ]; - let intermediate = PrecomputedMultimodalIntermediate { - modality: Modality::Video, - preprocessed, - images: vec![], - videos, - placeholders: vec![ + let intermediate = MultimodalIntermediate { + media: PreparedMedia::Videos(videos), + preprocessed: Arc::new(preprocessed), + structural_ranges: vec![ PlaceholderRange { offset: 30, length: 2, @@ -2271,7 +1753,7 @@ mod tests { length: 2, }, ], - patch_offsets: Some(vec![(30, 2), (40, 2)]), + feature_ranges: Some(vec![(30, 2), (40, 2)]), placeholder_token_id: Some(151656), field_layouts: HashMap::from([ ( @@ -2281,7 +1763,7 @@ mod tests { ("patches_per_video".to_string(), FieldLayout::Batched), ("video_grid_thw".to_string(), FieldLayout::Batched), ]), - keep_on_cpu_keys: vec![], + cpu_resident_tensor_keys: vec![], }; let assembled = assemble_tokenspeed(intermediate, None).unwrap(); @@ -2318,6 +1800,225 @@ mod tests { ); } + #[test] + #[cfg(unix)] + fn serialize_tokenspeed_encoder_inputs_packs_large_items_into_one_shm_segment() { + if !tokenspeed_shm_dev_writable() { + return; + } + + let first_values = vec![1.0f32; 16 * 1024]; + let second_values = vec![2.0f32; 16 * 1024]; + let first = ArrayD::from_shape_vec(IxDyn(&[16 * 1024]), first_values).unwrap(); + let second = ArrayD::from_shape_vec(IxDyn(&[16 * 1024]), second_values).unwrap(); + let first_view = first.view(); + let second_view = second.view(); + + let tensors = serialize_arrays_as_packed_tokenspeed_shm( + &[&first_view, &second_view], + "float32", + 1, + false, + ) + .expect("large encoder inputs should be packed"); + + let handles = tensors + .iter() + .map(|tensor| match &tensor.storage { + TokenSpeedTensorStorage::Shm(handle) => handle.clone(), + TokenSpeedTensorStorage::Inline(_) => panic!("expected packed SHM tensor"), + }) + .collect::>(); + + assert_eq!(handles.len(), 2); + assert_eq!(handles[0].name, handles[1].name); + assert_eq!(handles[0].offset, 0); + assert_eq!(handles[0].nbytes, (16 * 1024 * size_of::()) as u64); + assert_eq!(handles[1].offset, handles[0].nbytes); + assert_eq!(handles[1].nbytes, handles[0].nbytes); + + let mut file = fs::File::open(format!("/dev/shm/{}", handles[0].name)).unwrap(); + let mut bytes = [0u8; size_of::()]; + file.read_exact(&mut bytes).unwrap(); + assert_eq!(f32::from_le_bytes(bytes), 1.0); + file.seek(SeekFrom::Start(handles[1].offset)).unwrap(); + file.read_exact(&mut bytes).unwrap(); + assert_eq!(f32::from_le_bytes(bytes), 2.0); + + cleanup_tokenspeed_shm_handles(&handles); + } + + #[test] + #[cfg(unix)] + fn serialize_tokenspeed_encoder_inputs_packs_when_combined_size_reaches_threshold() { + if !tokenspeed_shm_dev_writable() { + return; + } + + let first = ArrayD::from_shape_vec(IxDyn(&[4]), vec![1.0f32; 4]).unwrap(); + let second = ArrayD::from_shape_vec(IxDyn(&[4]), vec![2.0f32; 4]).unwrap(); + let first_view = first.view(); + let second_view = second.view(); + let item_nbytes = (4 * size_of::()) as u64; + + let tensors = serialize_arrays_as_packed_tokenspeed_shm( + &[&first_view, &second_view], + "float32", + 32, + false, + ) + .expect("combined encoder input size should trigger packed SHM"); + + let handles = tensors + .iter() + .map(|tensor| match &tensor.storage { + TokenSpeedTensorStorage::Shm(handle) => handle.clone(), + TokenSpeedTensorStorage::Inline(_) => panic!("expected packed SHM tensor"), + }) + .collect::>(); + + assert_eq!(handles.len(), 2); + assert_eq!(handles[0].name, handles[1].name); + assert_eq!(handles[0].nbytes, item_nbytes); + assert_eq!(handles[1].offset, item_nbytes); + assert_eq!(handles[1].nbytes, item_nbytes); + + cleanup_tokenspeed_shm_handles(&handles); + } + + #[test] + #[cfg(unix)] + fn write_tokenspeed_shm_with_rejects_unexpected_byte_count() { + if !tokenspeed_shm_dev_writable() { + return; + } + + let err = write_tokenspeed_shm_with(4, |file| file.write_all(&[1, 2])) + .expect_err("short SHM writes must be rejected"); + + assert_eq!(err.kind(), std::io::ErrorKind::InvalidData); + } + + #[test] + fn serialize_tokenspeed_encoder_input_view_matches_owned_slice_bytes() { + let array = ArrayD::from_shape_vec( + IxDyn(&[3, 4]), + (0..12).map(|value| value as f32 + 0.25).collect(), + ) + .unwrap(); + let view = array.slice_axis(Axis(1), Slice::from(1..3)); + let owned = view.to_owned(); + + for dtype in ["float32", "bfloat16", "float16"] { + let (view_data, view_shape, view_dtype) = serialize_array_view_as_dtype(&view, dtype); + let (owned_data, owned_shape, owned_dtype) = + serialize_array_view_as_dtype(&owned.view(), dtype); + + assert_eq!(view_shape, owned_shape); + assert_eq!(view_dtype, owned_dtype); + assert_eq!(view_data, owned_data); + + let mut view_written = vec![0; view_data.len()]; + fill_array_as_dtype(&mut view_written, &view, dtype).unwrap(); + let mut owned_written = vec![0; owned_data.len()]; + fill_array_as_dtype(&mut owned_written, &owned.view(), dtype).unwrap(); + assert_eq!(view_written, owned_written); + } + } + + #[test] + fn serialize_tokenspeed_u16_inline_bytes_are_little_endian() { + let array = ArrayD::from_shape_vec(IxDyn(&[2]), vec![1.0_f32, -2.0_f32]).unwrap(); + + let (bf16, _, bf16_dtype) = serialize_array_view_as_dtype(&array.view(), "bfloat16"); + let (f16, _, f16_dtype) = serialize_array_view_as_dtype(&array.view(), "float16"); + + assert_eq!(bf16_dtype, "bfloat16"); + assert_eq!(bf16, vec![0x80, 0x3f, 0x00, 0xc0]); + assert_eq!(f16_dtype, "float16"); + assert_eq!(f16, vec![0x00, 0x3c, 0x00, 0xc0]); + } + + #[test] + fn serialize_tokenspeed_u16_parallel_matches_scalar_conversion() { + let values: Vec = (0..300_000) + .map(|index| (index as f32 - 150_000.0) / 257.0) + .collect(); + let array = ArrayD::from_shape_vec(IxDyn(&[values.len()]), values.clone()).unwrap(); + + for (dtype, convert) in [ + ("bfloat16", f32_to_bf16_bits as fn(f32) -> u16), + ("float16", f32_to_f16_bits as fn(f32) -> u16), + ] { + let (actual, _, _) = serialize_array_view_as_dtype(&array.view(), dtype); + let expected: Vec = values + .iter() + .flat_map(|&value| convert(value).to_le_bytes()) + .collect(); + assert_eq!(actual, expected); + + let mut written = vec![0; actual.len()]; + fill_array_as_dtype(&mut written, &array.view(), dtype).unwrap(); + assert_eq!(written, expected); + } + } + + #[test] + fn model_specific_tensor_bytes_are_little_endian() { + let float_tensor = + model_specific_to_tensor_bytes(&ModelSpecificValue::FloatVec(vec![1.0_f32, -2.0_f32])) + .unwrap(); + let int_tensor = + model_specific_to_tensor_bytes(&ModelSpecificValue::IntVec(vec![1_i64, -2_i64])) + .unwrap(); + let uint_tensor = + model_specific_to_tensor_bytes(&ModelSpecificValue::UintVec(vec![1_u32, 0x11223344])) + .unwrap(); + + assert_eq!( + float_tensor.data, + vec![0x00, 0x00, 0x80, 0x3f, 0x00, 0x00, 0x00, 0xc0] + ); + assert_eq!( + int_tensor.data, + vec![ + 0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xfe, 0xff, 0xff, 0xff, 0xff, 0xff, + 0xff, 0xff, + ] + ); + assert_eq!( + uint_tensor.data, + vec![0x01, 0x00, 0x00, 0x00, 0x44, 0x33, 0x22, 0x11] + ); + } + + #[test] + fn deferred_bf16_tokenspeed_tensor_matches_reference_conversion() { + let lut: [[f32; 256]; 3] = std::array::from_fn(|channel| { + std::array::from_fn(|value| value as f32 * (channel + 1) as f32 / 255.0) + }); + let raw = vec![0, 127, 255, 1, 128, 254]; + let deferred = + DeferredNormalizedEncoderInput::new(raw.clone(), vec![2, 3], lut, 1).unwrap(); + + let tensor = + serialize_deferred_bf16_tokenspeed_tensor(&deferred, false, usize::MAX, false).unwrap(); + let TokenSpeedTensorStorage::Inline(data) = tensor.storage else { + panic!("expected inline deferred tensor"); + }; + let expected = raw + .iter() + .enumerate() + .flat_map(|(index, &value)| { + f32_to_bf16_bits(lut[index % 3][value as usize]).to_le_bytes() + }) + .collect::>(); + + assert_eq!(data, expected); + assert_eq!(tensor.shape, vec![2, 3]); + assert_eq!(tensor.dtype, "bfloat16"); + } + // ------------------------------------------------------------------ // MultimodalConfigRegistry tests // ------------------------------------------------------------------ diff --git a/model_gateway/src/routers/grpc/multimodal/assembly.rs b/model_gateway/src/routers/grpc/multimodal/assembly.rs new file mode 100644 index 000000000..7110ed42c --- /dev/null +++ b/model_gateway/src/routers/grpc/multimodal/assembly.rs @@ -0,0 +1,148 @@ +// --------------------------------------------------------------------------- +// Assembly: convert MultimodalIntermediate → backend-specific MultimodalData +// --------------------------------------------------------------------------- + +mod serialization; +mod tokenspeed; + +use std::sync::Arc; + +use anyhow::Result; +use llm_multimodal::{ImageFrame, MultimodalRuntime, PreprocessedEncoderInputs}; +#[cfg(test)] +pub(super) use serialization::model_specific_to_tensor_bytes; +use serialization::{serialize_encoder_input, serialize_model_specific}; +pub(super) use tokenspeed::assemble_tokenspeed; +#[cfg(test)] +pub(super) use tokenspeed::{ + effective_tokenspeed_transport_mode, f32_to_bf16_bits, f32_to_f16_bits, fill_array_as_dtype, + flat_item_span, flat_item_spans, hash_hex_strings, local_shm_namespace_id, + placeholders_for_items, serialize_array_view_as_dtype, + serialize_arrays_as_packed_tokenspeed_shm, serialize_deferred_bf16_tokenspeed_tensor, + validate_tokenspeed_item_spans, +}; + +use super::{MultimodalIntermediate, PreparedMedia}; +use crate::routers::grpc::{ + client::GrpcClient, + context::WorkerSelection, + proto_wrapper::{SglangMultimodalData, TrtllmMultimodalData, VllmMultimodalData}, + MultimodalData, +}; + +/// Assemble backend-specific multimodal data from the intermediate. +/// +/// Called in request_building after worker selection, when the backend is known. +#[expect( + clippy::unreachable, + reason = "MLX multimodal rejected by caller before reaching here" +)] +pub(crate) fn assemble_multimodal_data( + intermediate: MultimodalIntermediate, + client: &GrpcClient, + workers: Option<&WorkerSelection>, + runtime: &MultimodalRuntime, +) -> Result { + runtime.run_cpu(|| match client { + GrpcClient::Sglang(_) => Ok(MultimodalData::Sglang(assemble_sglang( + materialize_encoder_input(intermediate)?, + )?)), + GrpcClient::Vllm(_) => Ok(MultimodalData::Vllm(assemble_vllm( + materialize_encoder_input(intermediate)?, + )?)), + GrpcClient::Trtllm(_) => Ok(MultimodalData::Trtllm(assemble_trtllm(intermediate)?)), + GrpcClient::TokenSpeed(_) => Ok(MultimodalData::TokenSpeed(assemble_tokenspeed( + intermediate, + workers, + )?)), + GrpcClient::Mlx(_) => unreachable!( + "caller rejects multimodal for MLX in build_chat_request/build_messages_request" + ), + }) +} + +fn materialize_encoder_input( + mut intermediate: MultimodalIntermediate, +) -> Result { + Arc::make_mut(&mut intermediate.preprocessed) + .materialize_encoder_input() + .map_err(|error| anyhow::anyhow!("failed to materialize encoder input: {error}"))?; + Ok(intermediate) +} + +fn assemble_sglang(intermediate: MultimodalIntermediate) -> Result { + let (pixel_values, pixel_values_shape) = serialize_encoder_input(&intermediate.preprocessed)?; + let model_specific_tensors = + serialize_model_specific(&intermediate.preprocessed.model_specific); + let image_data = prepared_images(&intermediate.media, "SGLang")? + .iter() + .map(|f| f.raw_bytes.to_vec()) + .collect(); + // Prefer encoder-feature ranges; fall back to full structural ranges. + let mm_placeholders = intermediate + .feature_ranges + .filter(|ranges| !ranges.is_empty()) + .unwrap_or_else(|| { + intermediate + .structural_ranges + .iter() + .map(|p| (p.offset as u32, p.length as u32)) + .collect() + }); + + Ok(SglangMultimodalData { + image_data, + pixel_values, + pixel_values_shape, + model_specific_tensors, + im_token_id: intermediate.placeholder_token_id, + mm_placeholders, + }) +} + +fn assemble_vllm(intermediate: MultimodalIntermediate) -> Result { + let (pixel_values, pixel_values_shape) = serialize_encoder_input(&intermediate.preprocessed)?; + let model_specific_tensors = + serialize_model_specific(&intermediate.preprocessed.model_specific); + let mm_hashes = prepared_images(&intermediate.media, "vLLM")? + .iter() + .map(|frame| frame.hash.clone()) + .collect(); + let mm_placeholders = intermediate + .structural_ranges + .iter() + .map(|p| (p.offset as u32, p.length as u32)) + .collect(); + let batched_keys = PreprocessedEncoderInputs::batched_keys(&intermediate.field_layouts); + let flat_keys = PreprocessedEncoderInputs::flat_keys(&intermediate.field_layouts); + + Ok(VllmMultimodalData { + pixel_values, + pixel_values_shape, + model_specific_tensors, + im_token_id: intermediate.placeholder_token_id, + mm_placeholders, + mm_hashes, + batched_keys, + flat_keys, + keep_on_cpu_keys: intermediate.cpu_resident_tensor_keys, + }) +} + +fn assemble_trtllm(intermediate: MultimodalIntermediate) -> Result { + let image_data = prepared_images(&intermediate.media, "TRT-LLM")? + .iter() + .map(|f| f.raw_bytes.to_vec()) + .collect(); + Ok(TrtllmMultimodalData { image_data }) +} + +fn prepared_images<'a>(media: &'a PreparedMedia, backend: &str) -> Result<&'a [Arc]> { + match media { + PreparedMedia::Images(images) => Ok(images), + PreparedMedia::Videos(_) => Err(anyhow::anyhow!( + "{backend} multimodal path currently supports image inputs only; got {}", + media.modality() + )), + } +} diff --git a/model_gateway/src/routers/grpc/multimodal/assembly/serialization.rs b/model_gateway/src/routers/grpc/multimodal/assembly/serialization.rs new file mode 100644 index 000000000..d9af4f5f3 --- /dev/null +++ b/model_gateway/src/routers/grpc/multimodal/assembly/serialization.rs @@ -0,0 +1,159 @@ +/// Serialize the primary encoder input ndarray to raw little-endian f32 bytes + shape. +use std::{collections::HashMap, mem::size_of}; + +use anyhow::Result; +use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs}; +use ndarray::ArrayD; +use tracing::warn; + +use crate::routers::grpc::TensorBytes; + +pub(super) fn serialize_encoder_input( + preprocessed: &PreprocessedEncoderInputs, +) -> Result<(Vec, Vec)> { + let encoder_input = preprocessed + .encoder_input + .dense() + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + Ok(serialize_array(encoder_input)) +} + +fn serialize_array(encoder_input: &ArrayD) -> (Vec, Vec) { + let encoder_bytes: Vec = if let Some(encoder_slice) = encoder_input + // Fast path only for C-contiguous arrays, whose memory order equals + // logical (row-major) order. A non-C-contiguous array (e.g. a + // Fortran-contiguous view) falls through to logical `.iter()` below; + // `as_slice_memory_order()` is deliberately NOT used as a fallback + // because it would serialize such arrays in the wrong dimension order. + .as_slice() + { + // Zero-copy reinterpret: &[f32] → &[u8] on little-endian (x86). + // This replaces the per-element flat_map(to_le_bytes) which was the + // #1 CPU hotspot (13% of SMG CPU in profiling). + #[cfg(target_endian = "little")] + { + let byte_slice: &[u8] = bytemuck::cast_slice(encoder_slice); + byte_slice.to_vec() + } + #[cfg(not(target_endian = "little"))] + { + f32_values_to_le_bytes(encoder_slice.iter().copied(), encoder_slice.len()) + } + } else { + // Non-C-contiguous array: `.iter()` walks in logical (row-major) order, + // which matches the shape. + f32_values_to_le_bytes(encoder_input.iter().copied(), encoder_input.len()) + }; + (encoder_bytes, array_shape(encoder_input)) +} + +fn array_shape(encoder_input: &ArrayD) -> Vec { + encoder_input.shape().iter().map(|&d| d as u32).collect() +} + +fn f32_values_to_le_bytes(values: I, len: usize) -> Vec +where + I: Iterator, +{ + let mut bytes = Vec::with_capacity(len * size_of::()); + for value in values { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes +} + +/// Serialize model-specific values to TensorBytes. +pub(super) fn serialize_model_specific( + model_specific: &HashMap, +) -> HashMap { + model_specific + .iter() + .filter_map(|(key, value)| match model_specific_to_tensor_bytes(value) { + Some(tensor) => Some((key.clone(), tensor)), + None => { + warn!(tensor_key = %key, "Dropping unsupported model_specific value during multimodal serialization"); + None + } + }) + .collect() +} + +/// Convert a model-specific value to backend-agnostic TensorBytes. +pub(in crate::routers::grpc::multimodal) fn model_specific_to_tensor_bytes( + value: &ModelSpecificValue, +) -> Option { + match value { + ModelSpecificValue::Tensor { data, shape } => Some(TensorBytes { + data: f32_slice_to_le_bytes(data), + shape: shape.iter().map(|&d| d as u32).collect(), + dtype: "float32".to_string(), + }), + ModelSpecificValue::IntTensor { data, shape } => Some(TensorBytes { + data: i64_slice_to_le_bytes(data), + shape: shape.iter().map(|&d| d as u32).collect(), + dtype: "int64".to_string(), + }), + ModelSpecificValue::UintTensor { data, shape } => Some(TensorBytes { + data: u32_slice_to_le_bytes(data), + shape: shape.iter().map(|&d| d as u32).collect(), + dtype: "uint32".to_string(), + }), + ModelSpecificValue::UintVec(v) => Some(TensorBytes { + data: u32_slice_to_le_bytes(v), + shape: vec![v.len() as u32], + dtype: "uint32".to_string(), + }), + ModelSpecificValue::IntVec(v) => Some(TensorBytes { + data: i64_slice_to_le_bytes(v), + shape: vec![v.len() as u32], + dtype: "int64".to_string(), + }), + ModelSpecificValue::FloatVec(v) => Some(TensorBytes { + data: f32_slice_to_le_bytes(v), + shape: vec![v.len() as u32], + dtype: "float32".to_string(), + }), + _ => None, + } +} + +fn f32_slice_to_le_bytes(values: &[f32]) -> Vec { + #[cfg(target_endian = "little")] + { + bytemuck::cast_slice(values).to_vec() + } + #[cfg(not(target_endian = "little"))] + { + f32_values_to_le_bytes(values.iter().copied(), values.len()) + } +} + +fn i64_slice_to_le_bytes(values: &[i64]) -> Vec { + #[cfg(target_endian = "little")] + { + bytemuck::cast_slice(values).to_vec() + } + #[cfg(not(target_endian = "little"))] + { + let mut bytes = Vec::with_capacity(values.len() * size_of::()); + for &value in values { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes + } +} + +fn u32_slice_to_le_bytes(values: &[u32]) -> Vec { + #[cfg(target_endian = "little")] + { + bytemuck::cast_slice(values).to_vec() + } + #[cfg(not(target_endian = "little"))] + { + let mut bytes = Vec::with_capacity(values.len() * size_of::()); + for &value in values { + bytes.extend_from_slice(&value.to_le_bytes()); + } + bytes + } +} diff --git a/model_gateway/src/routers/grpc/multimodal/assembly/tokenspeed.rs b/model_gateway/src/routers/grpc/multimodal/assembly/tokenspeed.rs new file mode 100644 index 000000000..8cdaad32c --- /dev/null +++ b/model_gateway/src/routers/grpc/multimodal/assembly/tokenspeed.rs @@ -0,0 +1,1193 @@ +mod transport; + +use std::{collections::HashMap, mem::size_of, sync::Arc, time::Instant}; + +use anyhow::{Context, Result}; +use llm_multimodal::{ + vision::transforms::preprocess_parallelism, DeferredNormalizedEncoderInput, FieldLayout, + Modality, ModelSpecificValue, PlaceholderRange, PreprocessedEncoderInputs, +}; +use ndarray::{ArrayD, ArrayViewD, Axis, Slice}; +use rayon::prelude::*; +use tracing::{info, warn}; +#[cfg(test)] +pub(in crate::routers::grpc::multimodal) use transport::{ + effective_tokenspeed_transport_mode, local_shm_namespace_id, +}; +use transport::{resolve_tokenspeed_shm_enabled, tokenspeed_encoder_input_dtype}; + +use super::{ + super::{log_mm_timing_enabled, MultimodalIntermediate}, + serialization::model_specific_to_tensor_bytes, +}; +use crate::routers::grpc::{ + context::WorkerSelection, + proto_wrapper::{ + tokenspeed_mm_shm_min_bytes, write_tokenspeed_shm_mapped, TensorBytes, TokenSpeedModality, + TokenSpeedMultimodalData, TokenSpeedMultimodalItem, TokenSpeedTensor, + }, +}; + +pub(in crate::routers::grpc::multimodal) fn assemble_tokenspeed( + mut intermediate: MultimodalIntermediate, + workers: Option<&WorkerSelection>, +) -> Result { + let log_timing = log_mm_timing_enabled(); + let total_started = log_timing.then(Instant::now); + let source_modality = intermediate.media.modality(); + // Resolve the multimodal tensor transport once per request: `shm` always on, + // `auto` only when the worker is verified to share /dev/shm (matching + // namespace token), otherwise inline. See `worker_shares_dev_shm`. + let shm_enabled = resolve_tokenspeed_shm_enabled(source_modality, workers); + // Prefer encoder-feature ranges; fall back to full structural ranges. + let encoder_input_dtype = tokenspeed_encoder_input_dtype(source_modality, workers); + let encoder_input_dtype = canonical_tokenspeed_encoder_dtype(&encoder_input_dtype); + if encoder_input_dtype != "bfloat16" && intermediate.preprocessed.encoder_input.is_deferred() { + Arc::make_mut(&mut intermediate.preprocessed) + .materialize_encoder_input() + .map_err(|error| anyhow::anyhow!("failed to materialize encoder input: {error}"))?; + } + let feature_ranges = intermediate + .feature_ranges + .as_deref() + .filter(|ranges| !ranges.is_empty()) + .unwrap_or(&[]); + + let modality = match source_modality { + Modality::Image => TokenSpeedModality::Image, + Modality::Video => TokenSpeedModality::Video, + Modality::Audio => TokenSpeedModality::Audio, + Modality::ImageEmbeds => TokenSpeedModality::Image, + }; + + let item_count = precomputed_multimodal_item_count(&intermediate)?; + let flat_spans = flat_item_spans( + &intermediate.preprocessed.model_specific, + &intermediate.field_layouts, + item_count, + )?; + validate_tokenspeed_item_spans( + intermediate.preprocessed.as_ref(), + &intermediate.field_layouts, + &flat_spans, + item_count, + )?; + let mm_placeholders_by_item = + placeholders_for_items(&intermediate.structural_ranges, feature_ranges); + anyhow::ensure!( + mm_placeholders_by_item.len() == item_count, + "precomputed multimodal assembly placeholder item count mismatch: modality={}, placeholder_item_count={}, item_count={item_count}", + source_modality, + mm_placeholders_by_item.len() + ); + let mut mm_placeholders_by_item = mm_placeholders_by_item.into_iter(); + if let Some(deferred) = intermediate + .preprocessed + .encoder_input + .deferred_normalized() + { + anyhow::ensure!( + item_count == 1, + "deferred TokenSpeed encoder input currently requires one multimodal item" + ); + // TODO: Add native typed payload support for vLLM/SGLang before using + // deferred BF16 outside TokenSpeed; converting BF16 back to FP32 would + // not preserve their current FP32 contract. + let model_specific_started = log_timing.then(Instant::now); + let model_specific_tensors = serialize_model_specific_for_item( + &intermediate.preprocessed.model_specific, + &intermediate.field_layouts, + &flat_spans, + 0, + )?; + let model_specific_serialize_ms = + model_specific_started.map(|started| started.elapsed().as_secs_f64() * 1000.0); + let mm_placeholders = mm_placeholders_by_item + .next() + .ok_or_else(|| anyhow::anyhow!("missing placeholders for multimodal item 0"))?; + let content_hash = content_hash_for_item(&intermediate, 0); + let encoder_input_started = log_timing.then(Instant::now); + let encoder_input = serialize_deferred_bf16_tokenspeed_tensor( + deferred, + shm_enabled, + tokenspeed_mm_shm_min_bytes(), + log_timing, + )?; + let encoder_input_serialize_ms = + encoder_input_started.map(|started| started.elapsed().as_secs_f64() * 1000.0); + if log_timing { + info!( + modality = ?modality, + item_index = 0, + encoder_input_dtype = %encoder_input.dtype, + encoder_input_bytes = encoder_input.nbytes(), + encoder_input_shape = ?encoder_input.shape, + model_specific_tensor_count = model_specific_tensors.len(), + encoder_input_serialize_ms = encoder_input_serialize_ms.unwrap_or_default(), + model_specific_serialize_ms = model_specific_serialize_ms.unwrap_or_default(), + "smg_mm_timing assemble_tokenspeed_item" + ); + } + if let Some(total_started) = total_started { + info!( + modality = ?modality, + item_count = 1, + total_ms = total_started.elapsed().as_secs_f64() * 1000.0, + "smg_mm_timing assemble_tokenspeed" + ); + } + return Ok(TokenSpeedMultimodalData { + items: vec![TokenSpeedMultimodalItem { + modality, + encoder_input, + model_specific_tensors, + placeholder_token_id: intermediate.placeholder_token_id, + mm_placeholders, + content_hash, + }], + shm_enabled, + }); + } + + let mut pending_items: Vec> = Vec::with_capacity(item_count); + for item_index in 0..item_count { + let item_encoder_input = encoder_input_for_item( + &intermediate.preprocessed, + &intermediate.field_layouts, + &flat_spans, + item_index, + )?; + let model_specific_started = log_timing.then(Instant::now); + let model_specific_tensors = serialize_model_specific_for_item( + &intermediate.preprocessed.model_specific, + &intermediate.field_layouts, + &flat_spans, + item_index, + )?; + let model_specific_serialize_ms = + model_specific_started.map(|started| started.elapsed().as_secs_f64() * 1000.0); + let mm_placeholders = mm_placeholders_by_item.next().ok_or_else(|| { + anyhow::anyhow!("missing placeholders for multimodal item {item_index}") + })?; + let content_hash = content_hash_for_item(&intermediate, item_index); + + pending_items.push(PendingTokenSpeedItem { + encoder_input: item_encoder_input, + model_specific_tensors, + mm_placeholders, + content_hash, + model_specific_serialize_ms, + }); + } + + let encoder_input_started = log_timing.then(Instant::now); + let encoder_inputs = if item_count == 1 { + let min_shm_bytes = tokenspeed_mm_shm_min_bytes(); + pending_items + .iter() + .map(|item| { + serialize_array_as_tokenspeed_tensor( + &item.encoder_input, + &encoder_input_dtype, + shm_enabled, + min_shm_bytes, + log_timing, + ) + }) + .collect() + } else { + serialize_arrays_as_tokenspeed_tensors( + pending_items.iter().map(|item| &item.encoder_input), + &encoder_input_dtype, + shm_enabled, + ) + }; + let encoder_input_serialize_ms = + encoder_input_started.map(|started| started.elapsed().as_secs_f64() * 1000.0); + + let mut items: Vec = Vec::with_capacity(item_count); + for (item_index, (pending, encoder_input)) in + pending_items.into_iter().zip(encoder_inputs).enumerate() + { + if log_timing { + info!( + modality = ?modality, + item_index, + encoder_input_dtype = %encoder_input.dtype, + encoder_input_bytes = encoder_input.nbytes(), + encoder_input_shape = ?encoder_input.shape, + model_specific_tensor_count = pending.model_specific_tensors.len(), + encoder_input_serialize_ms = encoder_input_serialize_ms.unwrap_or_default(), + model_specific_serialize_ms = pending + .model_specific_serialize_ms + .unwrap_or_default(), + "smg_mm_timing assemble_tokenspeed_item" + ); + } + + items.push(TokenSpeedMultimodalItem { + modality, + encoder_input, + model_specific_tensors: pending.model_specific_tensors, + placeholder_token_id: intermediate.placeholder_token_id, + mm_placeholders: pending.mm_placeholders, + content_hash: pending.content_hash, + }); + } + + if let Some(total_started) = total_started { + info!( + modality = ?modality, + item_count = items.len(), + total_ms = total_started.elapsed().as_secs_f64() * 1000.0, + "smg_mm_timing assemble_tokenspeed" + ); + } + + Ok(TokenSpeedMultimodalData { items, shm_enabled }) +} + +struct PendingTokenSpeedItem<'a> { + encoder_input: ArrayViewD<'a, f32>, + model_specific_tensors: HashMap, + mm_placeholders: Vec<(u32, u32)>, + content_hash: Vec, + model_specific_serialize_ms: Option, +} + +type FlatItemSpans = HashMap>; + +fn precomputed_multimodal_item_count(intermediate: &MultimodalIntermediate) -> Result { + let modality = intermediate.media.modality(); + let media_count = intermediate.media.item_count(); + let token_count = intermediate.preprocessed.feature_token_counts.len(); + let placeholder_count = intermediate.structural_ranges.len(); + let item_count = token_count.max(media_count).max(placeholder_count); + anyhow::ensure!( + item_count > 0, + "precomputed multimodal assembly requires at least one item" + ); + if media_count > 0 { + anyhow::ensure!( + media_count == item_count, + "precomputed multimodal assembly media count mismatch: modality={modality}, media_count={media_count}, item_count={item_count}" + ); + } + anyhow::ensure!( + token_count == item_count, + "precomputed multimodal assembly token count mismatch: modality={modality}, token_count={token_count}, item_count={item_count}" + ); + anyhow::ensure!( + placeholder_count == item_count, + "precomputed multimodal assembly placeholder count mismatch: modality={modality}, placeholder_count={placeholder_count}, item_count={item_count}" + ); + Ok(item_count) +} + +pub(in crate::routers::grpc::multimodal) fn flat_item_spans( + model_specific: &HashMap, + field_layouts: &HashMap, + item_count: usize, +) -> Result { + let mut spans_by_sizes_key = HashMap::new(); + for layout in field_layouts.values() { + let FieldLayout::Flat { sizes_key } = layout else { + continue; + }; + if spans_by_sizes_key.contains_key(sizes_key) { + continue; + } + + let sizes_value = model_specific + .get(sizes_key) + .ok_or_else(|| anyhow::anyhow!("missing flat sizes tensor {sizes_key}"))?; + spans_by_sizes_key.insert( + sizes_key.clone(), + item_spans_from_model_specific_sizes(sizes_key, sizes_value, item_count)?, + ); + } + Ok(spans_by_sizes_key) +} + +fn item_spans_from_model_specific_sizes( + sizes_key: &str, + value: &ModelSpecificValue, + item_count: usize, +) -> Result> { + let sizes_len = match value { + ModelSpecificValue::IntTensor { data, .. } => data.len(), + ModelSpecificValue::UintTensor { data, .. } => data.len(), + ModelSpecificValue::IntVec(values) => values.len(), + ModelSpecificValue::UintVec(values) => values.len(), + _ => anyhow::bail!("unsupported flat sizes value type"), + }; + anyhow::ensure!( + sizes_len == item_count, + "flat sizes tensor {sizes_key} length mismatch: sizes_len={sizes_len}, item_count={item_count}", + ); + + let mut spans = Vec::with_capacity(item_count); + let mut start = 0usize; + + match value { + ModelSpecificValue::IntTensor { data, .. } => { + for &len in data { + push_item_span_from_i64(&mut spans, &mut start, len)?; + } + } + ModelSpecificValue::UintTensor { data, .. } => { + for &len in data { + push_item_span(&mut spans, &mut start, len as usize)?; + } + } + ModelSpecificValue::IntVec(values) => { + for &len in values { + push_item_span_from_i64(&mut spans, &mut start, len)?; + } + } + ModelSpecificValue::UintVec(values) => { + for &len in values { + push_item_span(&mut spans, &mut start, len as usize)?; + } + } + _ => anyhow::bail!("unsupported flat sizes value type"), + } + Ok(spans) +} + +fn push_item_span_from_i64( + spans: &mut Vec<(usize, usize)>, + start: &mut usize, + len: i64, +) -> Result<()> { + let len = usize::try_from(len).context("negative flat size")?; + push_item_span(spans, start, len) +} + +fn push_item_span(spans: &mut Vec<(usize, usize)>, start: &mut usize, len: usize) -> Result<()> { + spans.push((*start, len)); + *start = (*start) + .checked_add(len) + .ok_or_else(|| anyhow::anyhow!("flat size offset overflow"))?; + Ok(()) +} + +pub(in crate::routers::grpc::multimodal) fn validate_tokenspeed_item_spans( + preprocessed: &PreprocessedEncoderInputs, + field_layouts: &HashMap, + flat_spans: &FlatItemSpans, + item_count: usize, +) -> Result<()> { + let encoder_shape = preprocessed.encoder_input_shape(); + let encoder_first_dim = *encoder_shape + .first() + .ok_or_else(|| anyhow::anyhow!("encoder_input tensor must have a first dimension"))?; + let encoder_layout = field_layouts + .get("pixel_values") + .unwrap_or(&FieldLayout::Batched); + validate_tokenspeed_layout_first_dim( + "pixel_values", + encoder_layout, + encoder_first_dim, + flat_spans, + item_count, + )?; + + for (key, value) in &preprocessed.model_specific { + let Some(layout) = field_layouts.get(key) else { + continue; + }; + let first_dim = model_specific_first_dim(key, value)?; + validate_tokenspeed_layout_first_dim(key, layout, first_dim, flat_spans, item_count)?; + } + + Ok(()) +} + +fn validate_tokenspeed_layout_first_dim( + tensor_key: &str, + layout: &FieldLayout, + first_dim: usize, + flat_spans: &FlatItemSpans, + item_count: usize, +) -> Result<()> { + match layout { + FieldLayout::Batched => { + anyhow::ensure!( + first_dim == item_count, + "batched tensor {tensor_key} first dimension mismatch: first_dim={first_dim}, item_count={item_count}" + ); + } + FieldLayout::Flat { sizes_key } => { + let spans = flat_spans.get(sizes_key).ok_or_else(|| { + anyhow::anyhow!("missing flat spans for sizes tensor {sizes_key}") + })?; + let span_total = spans.iter().try_fold(0usize, |acc, (_, len)| { + acc.checked_add(*len) + .ok_or_else(|| anyhow::anyhow!("flat span total overflow for {tensor_key}")) + })?; + anyhow::ensure!( + span_total == first_dim, + "flat tensor {tensor_key} first dimension mismatch: span_total={span_total}, first_dim={first_dim}, sizes_key={sizes_key}" + ); + } + } + Ok(()) +} + +fn model_specific_first_dim(key: &str, value: &ModelSpecificValue) -> Result { + match value { + ModelSpecificValue::Tensor { shape, .. } + | ModelSpecificValue::IntTensor { shape, .. } + | ModelSpecificValue::UintTensor { shape, .. } => shape.first().copied().ok_or_else(|| { + anyhow::anyhow!("model_specific tensor {key} must have a first dimension") + }), + ModelSpecificValue::IntVec(values) => Ok(values.len()), + ModelSpecificValue::UintVec(values) => Ok(values.len()), + ModelSpecificValue::FloatVec(values) => Ok(values.len()), + ModelSpecificValue::TupleVec(values) => Ok(values.len()), + ModelSpecificValue::Int(_) | ModelSpecificValue::Float(_) | ModelSpecificValue::Bool(_) => { + anyhow::bail!("model_specific value {key} has no first dimension") + } + } +} + +pub(in crate::routers::grpc::multimodal) fn flat_item_span( + flat_spans: &FlatItemSpans, + sizes_key: &str, + item_index: usize, +) -> Result<(usize, usize)> { + flat_spans + .get(sizes_key) + .and_then(|spans| spans.get(item_index)) + .copied() + .ok_or_else(|| { + anyhow::anyhow!("missing flat span for sizes tensor {sizes_key} item {item_index}") + }) +} + +fn encoder_input_for_item<'a>( + preprocessed: &'a PreprocessedEncoderInputs, + field_layouts: &HashMap, + flat_spans: &FlatItemSpans, + item_index: usize, +) -> Result> { + // The field layout key remains "pixel_values" because it is the established + // model vision input name. Internally this tensor is the modality encoder + // input we pass to TokenSpeed. + let layout = field_layouts + .get("pixel_values") + .unwrap_or(&FieldLayout::Batched); + let encoder_input = preprocessed + .encoder_input + .dense() + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + match layout { + FieldLayout::Batched => slice_array_axis0(encoder_input, item_index, 1), + FieldLayout::Flat { sizes_key } => { + let (start, len) = flat_item_span(flat_spans, sizes_key, item_index)?; + slice_array_axis0(encoder_input, start, len) + } + } +} + +fn serialize_model_specific_for_item( + model_specific: &HashMap, + field_layouts: &HashMap, + flat_spans: &FlatItemSpans, + item_index: usize, +) -> Result> { + let mut serialized = HashMap::with_capacity(model_specific.len()); + for (key, value) in model_specific { + let tensor = match field_layouts.get(key) { + Some(FieldLayout::Batched) => { + let item_value = value + .slice_first_dim(item_index, 1) + .with_context(|| format!("failed to slice model_specific tensor {key}"))?; + model_specific_to_tensor_bytes(&item_value) + } + Some(FieldLayout::Flat { sizes_key }) => { + let (start, len) = flat_item_span(flat_spans, sizes_key, item_index)?; + let item_value = value + .slice_first_dim(start, len) + .with_context(|| format!("failed to slice flat model_specific tensor {key}"))?; + model_specific_to_tensor_bytes(&item_value) + } + None => model_specific_to_tensor_bytes(value), + }; + if let Some(tensor) = tensor { + serialized.insert(key.clone(), tensor); + } else { + warn!(tensor_key = %key, "Dropping unsupported model_specific value during multimodal serialization"); + } + } + Ok(serialized) +} + +pub(in crate::routers::grpc::multimodal) fn placeholders_for_items( + placeholders: &[PlaceholderRange], + patch_offsets: &[(u32, u32)], +) -> Vec> { + if placeholders.len() == 1 { + return vec![placeholders_for_item(&placeholders[0], patch_offsets)]; + } + + if patch_offsets.is_empty() { + return placeholders + .iter() + .map(|placeholder| vec![full_placeholder_range(placeholder)]) + .collect(); + } + + if patch_offsets.len() == placeholders.len() { + let mut by_item = Vec::with_capacity(placeholders.len()); + let mut one_patch_run_per_item = true; + for (placeholder, &(offset, length)) in placeholders.iter().zip(patch_offsets) { + let start = placeholder.offset as u32; + let end = start + placeholder.length as u32; + if offset < start || offset.saturating_add(length) > end { + one_patch_run_per_item = false; + break; + } + by_item.push(vec![(offset, length)]); + } + if one_patch_run_per_item { + return by_item; + } + } + + if !placeholder_ranges_sorted(placeholders) || !patch_offsets_sorted(patch_offsets) { + return placeholders + .iter() + .map(|placeholder| placeholders_for_item(placeholder, patch_offsets)) + .collect(); + } + + let mut by_item = Vec::with_capacity(placeholders.len()); + let mut patch_idx = 0usize; + for placeholder in placeholders { + let start = placeholder.offset as u32; + let end = start + placeholder.length as u32; + while patch_idx < patch_offsets.len() && patch_offsets[patch_idx].0 < start { + patch_idx += 1; + } + + let mut item_patch_offsets = Vec::new(); + let mut scan_idx = patch_idx; + while scan_idx < patch_offsets.len() { + let (offset, length) = patch_offsets[scan_idx]; + if offset >= end { + break; + } + if offset >= start && offset.saturating_add(length) <= end { + item_patch_offsets.push((offset, length)); + } + scan_idx += 1; + } + patch_idx = scan_idx; + + if item_patch_offsets.is_empty() { + by_item.push(vec![(start, end - start)]); + } else { + by_item.push(item_patch_offsets); + } + } + by_item +} + +fn placeholders_for_item( + placeholder: &PlaceholderRange, + patch_offsets: &[(u32, u32)], +) -> Vec<(u32, u32)> { + let start = placeholder.offset as u32; + let end = start + placeholder.length as u32; + if patch_offsets.is_empty() { + return vec![(start, end - start)]; + } + if patch_offsets.len() == 1 { + let (offset, length) = patch_offsets[0]; + return if offset >= start && offset.saturating_add(length) <= end { + vec![(offset, length)] + } else { + vec![(start, end - start)] + }; + } + + let item_patch_offsets = patch_offsets + .iter() + .copied() + .filter(|(offset, length)| *offset >= start && offset.saturating_add(*length) <= end) + .collect::>(); + if item_patch_offsets.is_empty() { + vec![(start, end - start)] + } else { + item_patch_offsets + } +} + +fn full_placeholder_range(placeholder: &PlaceholderRange) -> (u32, u32) { + let start = placeholder.offset as u32; + (start, placeholder.length as u32) +} + +fn placeholder_ranges_sorted(placeholders: &[PlaceholderRange]) -> bool { + placeholders + .windows(2) + .all(|window| window[0].offset <= window[1].offset) +} + +fn patch_offsets_sorted(patch_offsets: &[(u32, u32)]) -> bool { + patch_offsets + .windows(2) + .all(|window| window[0].0 <= window[1].0) +} + +fn content_hash_for_item(intermediate: &MultimodalIntermediate, item_index: usize) -> Vec { + intermediate + .media + .content_hash(item_index) + .map(|hash| hash_hex_strings(std::iter::once(hash))) + .unwrap_or_default() +} + +fn slice_array_axis0(array: &ArrayD, start: usize, len: usize) -> Result> { + let end = start + .checked_add(len) + .ok_or_else(|| anyhow::anyhow!("array slice range overflow"))?; + let rows = array.shape().first().copied().unwrap_or(0); + anyhow::ensure!( + end <= rows, + "array first-dimension slice {start}..{end} exceeds {rows}" + ); + Ok(array.slice_axis(Axis(0), Slice::from(start..end))) +} + +pub(in crate::routers::grpc::multimodal) fn hash_hex_strings<'a>( + hashes: impl Iterator, +) -> Vec { + let mut hasher = blake3::Hasher::new(); + for hash in hashes { + hasher.update(hash.as_bytes()); + } + hasher.finalize().as_bytes().to_vec() +} + +// --------------------------------------------------------------------------- +// Serialization helpers +// --------------------------------------------------------------------------- + +pub(in crate::routers::grpc::multimodal) fn serialize_deferred_bf16_tokenspeed_tensor( + encoder_input: &DeferredNormalizedEncoderInput, + shm_enabled: bool, + min_shm_bytes: usize, + log_timing: bool, +) -> Result { + let nbytes = encoder_input + .len() + .checked_mul(size_of::()) + .ok_or_else(|| anyhow::anyhow!("deferred BF16 encoder input size overflow"))?; + let shape = encoder_input + .shape() + .iter() + .map(|&dimension| { + u32::try_from(dimension) + .map_err(|_| anyhow::anyhow!("encoder input dimension exceeds u32")) + }) + .collect::>>()?; + + if shm_enabled && nbytes >= min_shm_bytes { + let timing_started = log_timing.then(Instant::now); + match write_tokenspeed_shm_mapped(nbytes, |output| { + encoder_input + .fill_bf16_le_bytes(output) + .map_err(|error| std::io::Error::other(error.to_string())) + }) { + Ok(handle) => { + if log_timing { + info!( + nbytes, + elapsed_ms = timing_started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(), + "smg_mm_timing tokenspeed_shm_write_deferred_bf16" + ); + } + return Ok(TokenSpeedTensor::shm(handle, shape, "bfloat16".to_string())); + } + Err(error) => { + use crate::observability::metrics::Metrics; + warn!( + ?error, + nbytes, + "Failed to write deferred BF16 encoder input to SHM; falling back to inline transport" + ); + Metrics::record_mm_shm_write_failure("tokenspeed"); + } + } + } + + let mut data = vec![0; nbytes]; + encoder_input + .fill_bf16_le_bytes(&mut data) + .map_err(|error| anyhow::anyhow!(error.to_string()))?; + Ok(TokenSpeedTensor::inline( + data, + shape, + "bfloat16".to_string(), + )) +} + +/// Serialize encoder input to the requested wire dtype. +fn serialize_arrays_as_tokenspeed_tensors<'view, 'item>( + encoder_inputs: impl ExactSizeIterator>, + dtype: &str, + shm_enabled: bool, +) -> Vec +where + 'view: 'item, +{ + let min_shm_bytes = tokenspeed_mm_shm_min_bytes(); + let log_timing = log_mm_timing_enabled(); + let item_count = encoder_inputs.len(); + if shm_enabled && item_count >= 2 { + let encoder_inputs = encoder_inputs.collect::>(); + if let Some(tensors) = serialize_arrays_as_packed_tokenspeed_shm( + &encoder_inputs, + dtype, + min_shm_bytes, + log_timing, + ) { + return tensors; + } + return encoder_inputs + .iter() + .map(|&encoder_input| { + serialize_array_as_tokenspeed_tensor( + encoder_input, + dtype, + shm_enabled, + min_shm_bytes, + log_timing, + ) + }) + .collect(); + } + + encoder_inputs + .map(|encoder_input| { + serialize_array_as_tokenspeed_tensor( + encoder_input, + dtype, + shm_enabled, + min_shm_bytes, + log_timing, + ) + }) + .collect() +} + +pub(in crate::routers::grpc::multimodal) fn serialize_arrays_as_packed_tokenspeed_shm( + encoder_inputs: &[&ArrayViewD<'_, f32>], + dtype: &str, + min_bytes: usize, + log_timing: bool, +) -> Option> { + if encoder_inputs.len() < 2 { + return None; + } + + let dtype = canonical_tokenspeed_encoder_dtype(dtype); + let mut offsets = Vec::with_capacity(encoder_inputs.len()); + let mut nbytes_by_item = Vec::with_capacity(encoder_inputs.len()); + let mut shapes = Vec::with_capacity(encoder_inputs.len()); + let mut total_nbytes = 0usize; + for &encoder_input in encoder_inputs { + let nbytes = tokenspeed_encoder_input_nbytes(encoder_input, &dtype)?; + offsets.push(total_nbytes); + nbytes_by_item.push(nbytes); + shapes.push(array_view_shape(encoder_input)); + total_nbytes = total_nbytes.checked_add(nbytes)?; + } + if total_nbytes < min_bytes { + return None; + } + + let timing_started = log_timing.then(Instant::now); + match write_tokenspeed_shm_mapped(total_nbytes, |output| { + for ((&encoder_input, &offset), &nbytes) in + encoder_inputs.iter().zip(&offsets).zip(&nbytes_by_item) + { + fill_array_as_dtype(&mut output[offset..offset + nbytes], encoder_input, &dtype)?; + } + Ok(()) + }) { + Ok(base_handle) => { + if log_timing { + info!( + item_count = encoder_inputs.len(), + nbytes = total_nbytes, + elapsed_ms = timing_started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(), + "smg_mm_timing tokenspeed_shm_write_packed" + ); + } + + Some( + shapes + .into_iter() + .zip(offsets) + .zip(nbytes_by_item) + .map(|((shape, offset), nbytes)| { + let mut handle = base_handle.clone(); + handle.offset = offset as u64; + handle.nbytes = nbytes as u64; + TokenSpeedTensor::shm(handle, shape, dtype.clone()) + }) + .collect(), + ) + } + Err(error) => { + use crate::observability::metrics::Metrics; + warn!( + ?error, + item_count = encoder_inputs.len(), + nbytes = total_nbytes, + dtype = %dtype, + "Failed to write packed TokenSpeed encoder inputs to SHM; falling back to per-item transport" + ); + Metrics::record_mm_shm_write_failure("tokenspeed"); + None + } + } +} + +fn serialize_array_as_tokenspeed_tensor( + encoder_input: &ArrayViewD<'_, f32>, + dtype: &str, + shm_enabled: bool, + min_shm_bytes: usize, + log_timing: bool, +) -> TokenSpeedTensor { + let dtype = canonical_tokenspeed_encoder_dtype(dtype); + let shape = array_view_shape(encoder_input); + let Some(nbytes) = tokenspeed_encoder_input_nbytes(encoder_input, &dtype) else { + warn!( + dtype = %dtype, + shape = ?shape, + "TokenSpeed encoder input byte length overflow; falling back to inline serialization" + ); + let (data, shape, dtype) = serialize_array_view_as_dtype(encoder_input, &dtype); + return TokenSpeedTensor::inline(data, shape, dtype); + }; + + if shm_enabled && nbytes >= min_shm_bytes { + let timing_started = log_timing.then(Instant::now); + match write_tokenspeed_shm_mapped(nbytes, |output| { + fill_array_as_dtype(output, encoder_input, &dtype) + }) { + Ok(handle) => { + if log_timing { + info!( + nbytes, + elapsed_ms = timing_started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(), + "smg_mm_timing tokenspeed_shm_write_direct" + ); + } + return TokenSpeedTensor::shm(handle, shape, dtype); + } + Err(error) => { + use crate::observability::metrics::Metrics; + warn!( + ?error, + nbytes, + dtype = %dtype, + "Failed to write TokenSpeed encoder input directly to SHM; falling back to bytes path" + ); + Metrics::record_mm_shm_write_failure("tokenspeed"); + } + } + } + + let (data, shape, dtype) = serialize_array_view_as_dtype(encoder_input, &dtype); + TokenSpeedTensor::inline(data, shape, dtype) +} + +fn canonical_tokenspeed_encoder_dtype(dtype: &str) -> String { + match canonical_float_dtype(dtype).as_deref() { + Some("float32") => "float32".to_string(), + Some("bfloat16") => "bfloat16".to_string(), + Some("float16") => "float16".to_string(), + _ => { + warn!( + dtype, + "Unsupported TokenSpeed encoder input dtype; falling back to float32" + ); + "float32".to_string() + } + } +} + +fn tokenspeed_encoder_input_nbytes( + encoder_input: &ArrayViewD<'_, f32>, + dtype: &str, +) -> Option { + encoder_input + .len() + .checked_mul(tokenspeed_encoder_input_element_size(dtype)) +} + +fn tokenspeed_encoder_input_element_size(dtype: &str) -> usize { + if dtype == "bfloat16" || dtype == "float16" { + size_of::() + } else { + size_of::() + } +} + +pub(in crate::routers::grpc::multimodal) fn fill_array_as_dtype( + output: &mut [u8], + encoder_input: &ArrayViewD<'_, f32>, + dtype: &str, +) -> std::io::Result<()> { + let expected = tokenspeed_encoder_input_nbytes(encoder_input, dtype).ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "TokenSpeed encoder input byte length overflow", + ) + })?; + if output.len() != expected { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "TokenSpeed encoder output has an unexpected byte length", + )); + } + + match dtype { + "float32" => { + fill_array_as_f32_bytes(output, encoder_input); + Ok(()) + } + "bfloat16" => { + fill_array_as_u16_bytes(output, encoder_input, f32_to_bf16_bits); + Ok(()) + } + "float16" => { + fill_array_as_u16_bytes(output, encoder_input, f32_to_f16_bits); + Ok(()) + } + other => Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + format!("unsupported TokenSpeed encoder input dtype: {other}"), + )), + } +} + +fn fill_array_as_f32_bytes(output: &mut [u8], encoder_input: &ArrayViewD<'_, f32>) { + if let Some(encoder_slice) = encoder_input.as_slice() { + #[cfg(target_endian = "little")] + output.copy_from_slice(bytemuck::cast_slice(encoder_slice)); + #[cfg(not(target_endian = "little"))] + fill_f32_values_as_f32_bytes(output, encoder_slice.iter().copied()); + } else { + fill_f32_values_as_f32_bytes(output, encoder_input.iter().copied()); + } +} + +fn fill_f32_values_as_f32_bytes(output: &mut [u8], values: I) +where + I: IntoIterator, +{ + for (output, value) in output.chunks_exact_mut(size_of::()).zip(values) { + output.copy_from_slice(&value.to_le_bytes()); + } +} + +fn fill_array_as_u16_bytes(output: &mut [u8], encoder_input: &ArrayViewD<'_, f32>, convert: F) +where + F: Fn(f32) -> u16 + Copy + Send + Sync, +{ + if let Some(encoder_slice) = encoder_input.as_slice() { + fill_f32_slice_as_u16_bytes(output, encoder_slice, convert); + } else { + fill_f32_values_as_u16_bytes(output, encoder_input.iter().copied(), convert); + } +} + +pub(in crate::routers::grpc::multimodal) fn serialize_array_view_as_dtype( + encoder_input: &ArrayViewD<'_, f32>, + dtype: &str, +) -> (Vec, Vec, String) { + match canonical_float_dtype(dtype).as_deref() { + Some("float32") => { + let data = serialize_array_view_f32_bytes(encoder_input); + (data, array_view_shape(encoder_input), "float32".to_string()) + } + Some("bfloat16") => ( + serialize_array_view_as_u16_bytes(encoder_input, f32_to_bf16_bits), + array_view_shape(encoder_input), + "bfloat16".to_string(), + ), + Some("float16") => ( + serialize_array_view_as_u16_bytes(encoder_input, f32_to_f16_bits), + array_view_shape(encoder_input), + "float16".to_string(), + ), + _ => { + warn!( + dtype, + "Unsupported TokenSpeed encoder input dtype; falling back to float32" + ); + let data = serialize_array_view_f32_bytes(encoder_input); + (data, array_view_shape(encoder_input), "float32".to_string()) + } + } +} + +fn serialize_array_view_f32_bytes(encoder_input: &ArrayViewD<'_, f32>) -> Vec { + if let Some(encoder_slice) = encoder_input + // Fast path only for C-contiguous views, whose memory order equals + // logical (row-major) order. Non-C-contiguous views fall through to + // logical `.iter()` below, preserving the wire order. + .as_slice() + { + #[cfg(target_endian = "little")] + { + return bytemuck::cast_slice(encoder_slice).to_vec(); + } + #[cfg(not(target_endian = "little"))] + { + return f32_values_to_le_bytes(encoder_slice.iter().copied(), encoder_slice.len()); + } + } + + f32_values_to_le_bytes(encoder_input.iter().copied(), encoder_input.len()) +} + +fn f32_values_to_le_bytes(values: I, len: usize) -> Vec +where + I: IntoIterator, +{ + let mut bytes = Vec::with_capacity(len * size_of::()); + extend_f32_le_bytes(&mut bytes, values); + bytes +} + +fn extend_f32_le_bytes(bytes: &mut Vec, values: I) +where + I: IntoIterator, +{ + for value in values { + bytes.extend_from_slice(&value.to_le_bytes()); + } +} + +fn serialize_array_view_as_u16_bytes(encoder_input: &ArrayViewD<'_, f32>, convert: F) -> Vec +where + F: Fn(f32) -> u16 + Copy + Send + Sync, +{ + let element_count = encoder_input.len(); + let mut bytes = vec![0u8; element_count * size_of::()]; + + if let Some(encoder_slice) = encoder_input.as_slice() { + fill_f32_slice_as_u16_bytes(&mut bytes, encoder_slice, convert); + } else { + fill_f32_values_as_u16_bytes(&mut bytes, encoder_input.iter().copied(), convert); + } + bytes +} + +fn fill_f32_slice_as_u16_bytes(bytes: &mut [u8], values: &[f32], convert: F) +where + F: Fn(f32) -> u16 + Copy + Send + Sync, +{ + debug_assert_eq!(bytes.len(), values.len() * size_of::()); + let workers = preprocess_parallelism(bytes.len(), values.len()); + if workers <= 1 { + fill_f32_values_as_u16_bytes(bytes, values.iter().copied(), convert); + return; + } + + let chunk_values = values.len().div_ceil(workers); + bytes + .par_chunks_mut(chunk_values * size_of::()) + .zip(values.par_chunks(chunk_values)) + .for_each(|(output, values)| { + fill_f32_values_as_u16_bytes(output, values.iter().copied(), convert); + }); +} + +fn fill_f32_values_as_u16_bytes(bytes: &mut [u8], values: I, convert: F) +where + I: IntoIterator, + F: Fn(f32) -> u16 + Copy, +{ + for (output, value) in bytes.chunks_exact_mut(size_of::()).zip(values) { + output.copy_from_slice(&convert(value).to_le_bytes()); + } +} + +fn canonical_float_dtype(dtype: &str) -> Option { + match dtype.trim().to_ascii_lowercase().as_str() { + "float32" | "fp32" | "f32" => Some("float32".to_string()), + "bfloat16" | "bf16" => Some("bfloat16".to_string()), + "float16" | "fp16" | "f16" | "half" => Some("float16".to_string()), + _ => None, + } +} + +fn array_view_shape(encoder_input: &ArrayViewD<'_, f32>) -> Vec { + encoder_input.shape().iter().map(|&d| d as u32).collect() +} + +#[inline] +pub(in crate::routers::grpc::multimodal) fn f32_to_bf16_bits(value: f32) -> u16 { + let bits = value.to_bits(); + let lsb = (bits >> 16) & 1; + let rounding_bias = 0x7fff + lsb; + (bits.wrapping_add(rounding_bias) >> 16) as u16 +} + +#[inline] +pub(in crate::routers::grpc::multimodal) fn f32_to_f16_bits(value: f32) -> u16 { + let bits = value.to_bits(); + let sign = ((bits >> 16) & 0x8000) as u16; + let exp = ((bits >> 23) & 0xff) as i32; + let mant = bits & 0x7fffff; + + if exp == 0xff { + return if mant == 0 { + sign | 0x7c00 + } else { + sign | 0x7e00 + }; + } + + let half_exp = exp - 127 + 15; + if half_exp >= 0x1f { + return sign | 0x7c00; + } + if half_exp <= 0 { + if half_exp < -10 { + return sign; + } + let mantissa = mant | 0x800000; + let shift = (14 - half_exp) as u32; + let mut half_mant = (mantissa >> shift) as u16; + let round_bit = (mantissa >> (shift - 1)) & 1; + let sticky = mantissa & ((1u32 << (shift - 1)) - 1); + if round_bit != 0 && (sticky != 0 || (half_mant & 1) != 0) { + half_mant += 1; + } + return sign | half_mant; + } + + let mut half = sign | ((half_exp as u16) << 10) | ((mant >> 13) as u16); + let round = mant & 0x1fff; + if round > 0x1000 || (round == 0x1000 && (half & 1) != 0) { + half += 1; + } + half +} diff --git a/model_gateway/src/routers/grpc/multimodal/assembly/tokenspeed/transport.rs b/model_gateway/src/routers/grpc/multimodal/assembly/tokenspeed/transport.rs new file mode 100644 index 000000000..5f0e34b36 --- /dev/null +++ b/model_gateway/src/routers/grpc/multimodal/assembly/tokenspeed/transport.rs @@ -0,0 +1,185 @@ +use std::sync::OnceLock; + +use llm_multimodal::Modality; +use tracing::{info, warn}; + +use crate::routers::grpc::{ + context::WorkerSelection, + proto_wrapper::{ + tokenspeed_mm_shm_min_bytes, tokenspeed_mm_tensor_transport_mode, + tokenspeed_shm_dev_writable, + }, +}; + +pub(super) fn tokenspeed_encoder_input_dtype( + modality: Modality, + workers: Option<&WorkerSelection>, +) -> String { + if let Some(dtype) = tokenspeed_encoder_input_dtype_from_env(modality) { + return dtype; + } + if let Some(dtype) = tokenspeed_encoder_input_dtype_from_worker(workers) { + return dtype; + } + "float32".to_string() +} + +fn tokenspeed_encoder_input_dtype_from_env(modality: Modality) -> Option { + static IMAGE_DTYPE: OnceLock> = OnceLock::new(); + static VIDEO_DTYPE: OnceLock> = OnceLock::new(); + static AUDIO_DTYPE: OnceLock> = OnceLock::new(); + static DEFAULT_DTYPE: OnceLock> = OnceLock::new(); + + let modality_dtype = match modality { + Modality::Image | Modality::ImageEmbeds => { + cached_env_dtype(&IMAGE_DTYPE, "SMG_TOKENSPEED_IMAGE_ENCODER_INPUT_DTYPE") + } + Modality::Video => { + cached_env_dtype(&VIDEO_DTYPE, "SMG_TOKENSPEED_VIDEO_ENCODER_INPUT_DTYPE") + } + Modality::Audio => { + cached_env_dtype(&AUDIO_DTYPE, "SMG_TOKENSPEED_AUDIO_ENCODER_INPUT_DTYPE") + } + }; + modality_dtype + .or_else(|| cached_env_dtype(&DEFAULT_DTYPE, "SMG_TOKENSPEED_ENCODER_INPUT_DTYPE")) +} + +fn cached_env_dtype(cell: &'static OnceLock>, name: &str) -> Option { + cell.get_or_init(|| std::env::var(name).ok().filter(|dtype| !dtype.is_empty())) + .clone() +} + +fn tokenspeed_encoder_input_dtype_from_worker(workers: Option<&WorkerSelection>) -> Option { + let worker = match workers? { + WorkerSelection::Single { worker } => worker, + WorkerSelection::Dual { prefill, .. } => prefill, + }; + worker + .metadata() + .spec + .labels + .get("multimodal_encoder_dtype") + .filter(|dtype| !dtype.is_empty()) + .cloned() +} + +/// Resolve whether large multimodal tensors should use the SHM transport for +/// this request. `shm` = always (legacy explicit opt-in); `auto` = only when the +/// worker is known to share SMG's `/dev/shm`; anything else (including unset or +/// `inline`) keeps the inline gRPC path. +pub(super) fn resolve_tokenspeed_shm_enabled( + modality: Modality, + workers: Option<&WorkerSelection>, +) -> bool { + let configured_mode = tokenspeed_mm_tensor_transport_mode(); + let mode = effective_tokenspeed_transport_mode(modality, &configured_mode); + log_tokenspeed_transport_config_once(&configured_mode, &mode, modality); + match mode.as_str() { + // SHM only ever happens when SMG can actually write /dev/shm. + "shm" => tokenspeed_shm_dev_writable(), + "auto" => worker_shares_dev_shm(workers) && tokenspeed_shm_dev_writable(), + "" | "inline" => false, + other => { + log_unknown_tokenspeed_transport_once(other); + false + } + } +} + +pub(in crate::routers::grpc::multimodal) fn effective_tokenspeed_transport_mode( + modality: Modality, + configured_mode: &str, +) -> String { + if !configured_mode.is_empty() { + return configured_mode.to_string(); + } + + match modality { + Modality::Video => "auto".to_string(), + Modality::Image | Modality::ImageEmbeds | Modality::Audio => "inline".to_string(), + } +} + +fn log_tokenspeed_transport_config_once( + configured_mode: &str, + effective_mode: &str, + modality: Modality, +) { + static LOGGED: OnceLock<()> = OnceLock::new(); + LOGGED.get_or_init(|| { + info!( + configured_mode, + effective_mode, + ?modality, + shm_min_bytes = tokenspeed_mm_shm_min_bytes(), + dev_writable = tokenspeed_shm_dev_writable(), + "TokenSpeed multimodal tensor transport configured" + ); + }); +} + +fn log_unknown_tokenspeed_transport_once(value: &str) { + static WARNED: OnceLock<()> = OnceLock::new(); + WARNED.get_or_init(|| { + warn!( + value, + "Unknown SMG_TOKENSPEED_MM_TENSOR_TRANSPORT value; expected inline|shm|auto, using inline" + ); + }); +} + +/// Whether the worker is *verified* to share SMG's `/dev/shm`, making the SHM +/// transport safe under `auto`. +/// +/// Rather than inferring locality from the worker URL (TCP loopback proves only +/// network locality, not a shared `/dev/shm`), the worker advertises its +/// `/dev/shm` filesystem identity (`:`) via +/// `GetServerInfo`, which discovery stores in the worker's `shm_namespace_id` +/// label. Two processes share `/dev/shm` iff these tokens match: `boot_id` pins +/// the host, and `st_dev` is the tmpfs superblock device, identical whenever the +/// same tmpfs backs both `/dev/shm` mounts — including separate containers that +/// share it via `--ipc`/bind-mount (where mount-namespace inodes differ but the +/// underlying superblock is the same). We compare the worker's token to ours: +/// equal ⇒ shared. A missing/empty token or any mismatch is treated as +/// non-sharing, so `auto` safely falls back to inline. +fn worker_shares_dev_shm(workers: Option<&WorkerSelection>) -> bool { + let Some(local) = local_shm_namespace_id() else { + return false; + }; + let worker = match workers { + Some(WorkerSelection::Single { worker }) => worker, + Some(WorkerSelection::Dual { prefill, .. }) => prefill, + None => return false, + }; + worker + .metadata() + .spec + .labels + .get("shm_namespace_id") + .is_some_and(|id| !id.is_empty() && id == local) +} + +/// This process's `/dev/shm` filesystem identity: `:`. +/// `boot_id` pins the host (it is not namespaced) and `st_dev` is the tmpfs +/// superblock device backing `/dev/shm`; together they identify the tmpfs so two +/// processes sharing it (even across containers via `--ipc`/bind-mount) produce +/// the same token. Computed once; `None` if it can't be determined (then `auto` +/// stays inline). +pub(in crate::routers::grpc::multimodal) fn local_shm_namespace_id() -> Option<&'static str> { + static ID: OnceLock> = OnceLock::new(); + ID.get_or_init(compute_shm_namespace_id).as_deref() +} + +#[cfg(unix)] +fn compute_shm_namespace_id() -> Option { + use std::os::unix::fs::MetadataExt; + let boot_id = std::fs::read_to_string("/proc/sys/kernel/random/boot_id").ok()?; + let shm_dev = std::fs::metadata("/dev/shm").ok()?.dev(); + Some(format!("{}:{shm_dev}", boot_id.trim())) +} + +#[cfg(not(unix))] +fn compute_shm_namespace_id() -> Option { + None +} diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index b591aabbe..1c929df1f 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -360,13 +360,15 @@ fn tokenspeed_tensor_payload(data: Vec, shm_enabled: bool) -> tokenspeed::te return tokenspeed::tensor_data::Payload::Inline(data); } - let started = Instant::now(); + let started = log_timing.then(Instant::now); match write_tokenspeed_shm(&data) { Ok(handle) => { if log_timing { tracing::info!( nbytes, - elapsed_ms = started.elapsed().as_secs_f64() * 1000.0, + elapsed_ms = started + .map(|started| started.elapsed().as_secs_f64() * 1000.0) + .unwrap_or_default(), "smg_mm_timing tokenspeed_shm_write" ); } @@ -387,9 +389,12 @@ fn tokenspeed_tensor_payload(data: Vec, shm_enabled: bool) -> tokenspeed::te } fn log_tokenspeed_mm_timing_enabled() -> bool { - std::env::var("SMG_LOG_MM_TIMING") - .map(|value| matches!(value.to_ascii_lowercase().as_str(), "1" | "true" | "yes")) - .unwrap_or(false) + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + std::env::var("SMG_LOG_MM_TIMING") + .map(|value| matches!(value.to_ascii_lowercase().as_str(), "1" | "true" | "yes")) + .unwrap_or(false) + }) } /// Multimodal tensor transport mode for the TokenSpeed backend. @@ -398,19 +403,26 @@ fn log_tokenspeed_mm_timing_enabled() -> bool { /// model-specific tensors); prompt `input_ids` are always sent inline. Set via /// `SMG_TOKENSPEED_MM_TENSOR_TRANSPORT`. pub fn tokenspeed_mm_tensor_transport_mode() -> String { - std::env::var("SMG_TOKENSPEED_MM_TENSOR_TRANSPORT") - .unwrap_or_default() - .trim() - .to_ascii_lowercase() + static MODE: OnceLock = OnceLock::new(); + MODE.get_or_init(|| { + std::env::var("SMG_TOKENSPEED_MM_TENSOR_TRANSPORT") + .unwrap_or_default() + .trim() + .to_ascii_lowercase() + }) + .clone() } /// Minimum multimodal tensor size (bytes) before the SHM transport is used. /// Set via `SMG_TOKENSPEED_MM_SHM_MIN_BYTES`. Defaults to 64 KiB. pub fn tokenspeed_mm_shm_min_bytes() -> usize { - std::env::var("SMG_TOKENSPEED_MM_SHM_MIN_BYTES") - .ok() - .and_then(|value| value.parse::().ok()) - .unwrap_or(64 * 1024) + static MIN_BYTES: OnceLock = OnceLock::new(); + *MIN_BYTES.get_or_init(|| { + std::env::var("SMG_TOKENSPEED_MM_SHM_MIN_BYTES") + .ok() + .and_then(|value| value.parse::().ok()) + .unwrap_or(64 * 1024) + }) } static TOKENSPEED_SHM_COUNTER: AtomicU64 = AtomicU64::new(0); @@ -419,6 +431,41 @@ fn write_tokenspeed_shm(data: &[u8]) -> std::io::Result { write_tokenspeed_shm_with(data.len(), |file| file.write_all(data)) } +pub(crate) struct CountingWriter { + inner: W, + bytes_written: usize, +} + +impl CountingWriter { + fn new(inner: W) -> Self { + Self { + inner, + bytes_written: 0, + } + } + + fn bytes_written(&self) -> usize { + self.bytes_written + } +} + +impl Write for CountingWriter { + fn write(&mut self, buf: &[u8]) -> std::io::Result { + let n = self.inner.write(buf)?; + self.bytes_written = self.bytes_written.checked_add(n).ok_or_else(|| { + std::io::Error::new( + std::io::ErrorKind::InvalidData, + "TokenSpeed SHM writer byte count overflowed", + ) + })?; + Ok(n) + } + + fn flush(&mut self) -> std::io::Result<()> { + self.inner.flush() + } +} + /// Whether SMG can actually create+write files under `/dev/shm`. Probed once; /// when false the SHM transport cannot work, so `auto`/`shm` must stay inline. pub fn tokenspeed_shm_dev_writable() -> bool { @@ -494,15 +541,14 @@ fn sweep_orphan_tokenspeed_shm_once() { }); } -// TODO: pack all of a request's tensors (encoder_input + model_specific) into -// ONE /dev/shm segment at running offsets instead of one file per tensor -// (ShmHandle.offset already exists, always 0 here). Needs consumer -// ShmTensorHandle offset support + a per-segment refcount so the segment is -// unlinked exactly once after all its tensors are consumed. Cleanliness / fewer -// files, not a measured speed win (tmpfs makes per-file syscalls negligible). -pub fn write_tokenspeed_shm_with( +// Multimodal assembly can pack multiple encoder_input tensors into one segment +// by cloning this returned handle and setting per-tensor offsets. Do not pack +// model-specific tensors into that same segment until the worker has segment +// refcounting: those tensors are read and unlinked by the gRPC servicer before +// encoder_input handles are materialized by TokenSpeed's multimodal planner. +pub(crate) fn write_tokenspeed_shm_with( nbytes: usize, - write_fn: impl FnOnce(&mut BufWriter) -> std::io::Result<()>, + write_fn: impl FnOnce(&mut CountingWriter>) -> std::io::Result<()>, ) -> std::io::Result { sweep_orphan_tokenspeed_shm_once(); let name = next_tokenspeed_shm_name(); @@ -516,14 +562,14 @@ pub fn write_tokenspeed_shm_with( opts.mode(0o600); } let file = opts.open(&path)?; - let mut writer = BufWriter::new(file); + let mut writer = CountingWriter::new(BufWriter::new(file)); if let Err(error) = write_fn(&mut writer) { drop(writer); let _ = remove_file(&path); return Err(error); } if let Err(error) = writer.flush().and_then(|()| { - if writer.get_ref().metadata()?.len() != nbytes as u64 { + if writer.bytes_written() != nbytes { return Err(std::io::Error::new( std::io::ErrorKind::InvalidData, "TokenSpeed SHM writer produced an unexpected byte length", @@ -544,6 +590,104 @@ pub fn write_tokenspeed_shm_with( }) } +/// Creates a fixed-size TokenSpeed SHM segment and exposes its mapped payload. +/// +/// This avoids an intermediate conversion buffer for large encoder tensors: +/// callers can convert directly into the pages the worker will consume. +pub(crate) fn write_tokenspeed_shm_mapped( + nbytes: usize, + write_fn: impl FnOnce(&mut [u8]) -> std::io::Result<()>, +) -> std::io::Result { + if nbytes == 0 { + return Err(std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "TokenSpeed SHM mapping must not be empty", + )); + } + + sweep_orphan_tokenspeed_shm_once(); + let name = next_tokenspeed_shm_name(); + let path = tokenspeed_shm_path(&name); + let mut opts = OpenOptions::new(); + opts.read(true).write(true).create_new(true); + #[cfg(unix)] + { + use std::os::unix::fs::OpenOptionsExt; + opts.mode(0o600); + } + let file = opts.open(&path)?; + let file_len = u64::try_from(nbytes).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "TokenSpeed SHM mapping length does not fit in u64", + ) + })?; + if let Err(error) = file.set_len(file_len) { + drop(file); + let _ = remove_file(&path); + return Err(error); + } + if let Err(error) = reserve_tokenspeed_shm(&file, nbytes) { + drop(file); + let _ = remove_file(&path); + return Err(error); + } + + let mut mapping = match map_tokenspeed_shm(&file, nbytes) { + Ok(mapping) => mapping, + Err(error) => { + drop(file); + let _ = remove_file(&path); + return Err(error); + } + }; + if let Err(error) = write_fn(&mut mapping) { + drop(mapping); + drop(file); + let _ = remove_file(&path); + return Err(error); + } + drop(mapping); + + Ok(tokenspeed::ShmHandle { + name, + offset: 0, + nbytes: file_len, + owner_id: format!("smg:{}", process::id()), + }) +} + +// Mapping a newly-created, exclusively owned file is the only unsafe operation +// needed for direct SHM serialization. The mapping cannot outlive `file` here. +#[expect(unsafe_code, reason = "memmap2 requires unsafe mapping creation")] +fn map_tokenspeed_shm(file: &std::fs::File, nbytes: usize) -> std::io::Result { + unsafe { memmap2::MmapOptions::new().len(nbytes).map_mut(file) } +} + +#[cfg(target_os = "linux")] +#[expect(unsafe_code, reason = "posix_fallocate is exposed through libc")] +fn reserve_tokenspeed_shm(file: &std::fs::File, nbytes: usize) -> std::io::Result<()> { + use std::os::fd::AsRawFd; + + let length = libc::off_t::try_from(nbytes).map_err(|_| { + std::io::Error::new( + std::io::ErrorKind::InvalidInput, + "TokenSpeed SHM reservation length does not fit in off_t", + ) + })?; + let result = unsafe { libc::posix_fallocate(file.as_raw_fd(), 0, length) }; + if result == 0 { + Ok(()) + } else { + Err(std::io::Error::from_raw_os_error(result)) + } +} + +#[cfg(not(target_os = "linux"))] +fn reserve_tokenspeed_shm(_file: &std::fs::File, _nbytes: usize) -> std::io::Result<()> { + Ok(()) +} + pub fn collect_tokenspeed_multimodal_inputs_shm_handles( inputs: &tokenspeed::MultimodalInputs, ) -> Vec { @@ -615,35 +759,6 @@ pub(crate) fn finish_tokenspeed_request( } } -/// Unlink the `/dev/shm` segments backing the encoder inputs of intermediate -/// `items` (plus an optional just-built `pending` tensor that hasn't been pushed -/// yet). Used when multimodal assembly aborts partway: the successfully built -/// `TokenSpeedTensor::Shm` segments would otherwise be dropped without their -/// handles ever reaching the send-path cleanup hooks, leaking files until the -/// next process sweep. Only the encoder input uses SHM (model-specific tensors -/// stay inline). MUST run on the error path only — the success path keeps the -/// files alive for the worker and unlinks them after the RPC. -pub(crate) fn cleanup_tokenspeed_items_encoder_shm( - items: &[TokenSpeedMultimodalItem], - pending: Option<&TokenSpeedTensor>, -) { - let mut handles = Vec::new(); - let mut push = |tensor: &TokenSpeedTensor| { - if let TokenSpeedTensorStorage::Shm(handle) = &tensor.storage { - handles.push(handle.clone()); - } - }; - for item in items { - push(&item.encoder_input); - } - if let Some(tensor) = pending { - push(tensor); - } - if !handles.is_empty() { - cleanup_tokenspeed_shm_handles(&handles); - } -} - fn collect_optional_tokenspeed_tensor_shm_handles( tensor: Option<&tokenspeed::TensorData>, handles: &mut Vec, diff --git a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs index e854fde6d..0870d4831 100644 --- a/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/chat/request_building.rs @@ -89,7 +89,17 @@ impl PipelineStage for ChatRequestBuildingStage { let multimodal_data = processed_messages .multimodal_intermediate .map(|intermediate| { - assemble_multimodal_data(intermediate, builder_client, ctx.state.workers.as_ref()) + let runtime = ctx + .components + .multimodal + .as_ref() + .ok_or_else(|| anyhow::anyhow!("multimodal runtime is unavailable"))?; + assemble_multimodal_data( + intermediate, + builder_client, + ctx.state.workers.as_ref(), + &runtime.runtime, + ) }) .transpose() .map_err(|e| { diff --git a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs index 4bef07971..4f82726ef 100644 --- a/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs +++ b/model_gateway/src/routers/grpc/regular/stages/messages/request_building.rs @@ -91,7 +91,17 @@ impl PipelineStage for MessageRequestBuildingStage { let multimodal_data = processed_messages .multimodal_intermediate .map(|intermediate| { - assemble_multimodal_data(intermediate, builder_client, ctx.state.workers.as_ref()) + let runtime = ctx + .components + .multimodal + .as_ref() + .ok_or_else(|| anyhow::anyhow!("multimodal runtime is unavailable"))?; + assemble_multimodal_data( + intermediate, + builder_client, + ctx.state.workers.as_ref(), + &runtime.runtime, + ) }) .transpose() .map_err(|e| {