feat(multimodal): optimize EPD encode routing

bindings/python/src/lib.rs

@@ -540,7 +540,7 @@ impl Router {
                     config::ConfigError::InvalidValue {
                         field: "multimodal_tensor_transport".to_string(),
                         value: value.to_string(),
-                        reason: "expected 'inline', 'shm', or 'auto'".to_string(),
+                        reason: "expected 'inline', 'shm', 'auto', or 'rdma'".to_string(),
                     }
                 })
             })

bindings/python/src/smg/router_args.py

@@ -571,9 +571,9 @@ def add_cli_args(
         parser.add_argument(
             f"--{prefix}multimodal-tensor-transport",
             type=str,
-            choices=["inline", "shm", "auto"],
+            choices=["inline", "shm", "auto", "rdma"],
             default=RouterArgs.multimodal_tensor_transport,
-            help="Multimodal tensor transport mode: inline (default), shm, or auto",
+            help="Multimodal tensor transport: inline (default), shm, auto, or rdma (NIXL lane; needs mm-rdma build)",
         )
         parser.add_argument(
             f"--{prefix}multimodal-shm-min-bytes",

crates/multimodal/src/vision/mod.rs

@@ -33,6 +33,7 @@ pub(crate) mod execution;
 pub mod preprocessor_config;
 pub mod processor;
 pub mod processors;
+pub(crate) mod scratch;
 pub mod transforms;
 
 // Re-export commonly used types

crates/multimodal/src/vision/processors/kimi_k25.rs

@@ -18,6 +18,7 @@ use ndarray::Array3;
 use crate::vision::{
     preprocessor_config::PreProcessorConfig,
     processor::{ModelSpecificValue, PreprocessedEncoderInputs, VisionPreProcessor},
+    scratch,
     transforms::{self, TransformError},
 };
 
@@ -156,7 +157,9 @@ impl KimiK25Processor {
         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 mut data = vec![0.0f32; 3 * canvas_pixels];
+        // Pooled: this per-image CHW buffer (tens of MB) is recycled by the
+        // caller after patch extraction, keeping its pages mapped and hot.
+        let mut data = scratch::take_f32(3 * canvas_pixels);
         let (r_plane, rest) = data.split_at_mut(canvas_pixels);
         let (g_plane, b_plane) = rest.split_at_mut(canvas_pixels);
 
@@ -193,17 +196,15 @@ impl KimiK25Processor {
     ///
     /// Uses row-based `copy_from_slice` instead of per-element indexing so the
     /// compiler can auto-vectorize the inner copy.
-    fn extract_patches(tensor: &Array3<f32>, patch_size: usize) -> Vec<f32> {
+    /// Append this image's patches directly into `out` (no per-image intermediate
+    /// Vec): `out` is the pooled batch buffer pre-sized for the whole request.
+    fn extract_patches_into(tensor: &Array3<f32>, patch_size: usize, out: &mut Vec<f32>) {
         let channels = tensor.shape()[0];
         let height = tensor.shape()[1];
         let width = tensor.shape()[2];
 
         let grid_h = height / patch_size;
         let grid_w = width / patch_size;
-        let num_patches = grid_h * grid_w;
-        let patch_features = channels * patch_size * patch_size;
-
-        let mut patches = Vec::with_capacity(num_patches * patch_features);
 
         // Get contiguous slice for direct row addressing
         let flat = tensor.as_standard_layout();
@@ -223,13 +224,11 @@ impl KimiK25Processor {
                     let plane_offset = c * height * width;
                     for ph in 0..patch_size {
                         let row_start = plane_offset + (h_start + ph) * width + w_start;
-                        patches.extend_from_slice(&data[row_start..row_start + patch_size]);
+                        out.extend_from_slice(&data[row_start..row_start + patch_size]);
                     }
                 }
             }
         }
-
-        patches
     }
 }
 
@@ -255,7 +254,18 @@ impl VisionPreProcessor for KimiK25Processor {
         let mean = config.get_image_mean();
         let std = config.get_image_std();
 
-        let mut all_patches: Vec<f32> = Vec::new();
+        // Pre-size the pooled batch buffer exactly (patch_features per patch =
+        // 3 * patch_size^2; this is the data plane's hottest allocation).
+        let patch_features = 3 * self.patch_size * self.patch_size;
+        let mut estimated_total = 0usize;
+        for image in images {
+            let (w, h) = image.dimensions();
+            let cfg = self.compute_resize_config(w as usize, h as usize);
+            let grid_h = (cfg.new_height + cfg.pad_height) / self.patch_size;
+            let grid_w = (cfg.new_width + cfg.pad_width) / self.patch_size;
+            estimated_total += grid_h * grid_w * patch_features;
+        }
+        let mut all_patches: Vec<f32> = scratch::take_f32_cap(estimated_total);
         let mut patches_per_image: Vec<i64> = 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());
@@ -280,8 +290,11 @@ impl VisionPreProcessor for KimiK25Processor {
             let num_patches = grid_h * grid_w;
             feature_token_counts.push(cfg.num_tokens);
 
-            let patches = Self::extract_patches(&tensor, self.patch_size);
-            all_patches.extend(patches);
+            // Patchify directly into the pooled batch buffer, then recycle the
+            // CHW tensor's storage (standard layout, offset 0) for the next image.
+            Self::extract_patches_into(&tensor, self.patch_size, &mut all_patches);
+            let (storage, _offset) = tensor.into_raw_vec_and_offset();
+            scratch::give_f32(storage);
             patches_per_image.push(num_patches as i64);
         }
 

crates/multimodal/src/vision/processors/qwen3_vl.rs

@@ -452,6 +452,74 @@ mod tests {
         }
     }
 
+    // Per-image-independence guard for the smg gateway pixel cache (EPD P1): the
+    // cache stores one entry per image and reassembles requests from per-image
+    // payloads, which is only sound if preprocessing image i is independent of the
+    // other images in the batch. This asserts that a batched preprocess equals the
+    // concatenation of per-image preprocesses (encoder_input rows, grid_thw rows,
+    // and feature_token_counts). If a processor ever introduces cross-image work,
+    // this fails and the cache assumption must be revisited.
+    #[test]
+    fn per_image_preprocess_equals_batched_slices() {
+        use ndarray::{Axis, Slice};
+
+        let processor = Qwen3VLProcessor::new();
+        let config = PreProcessorConfig {
+            image_mean: Some(QWEN3_MEAN.to_vec()),
+            image_std: Some(QWEN3_STD.to_vec()),
+            ..Default::default()
+        };
+
+        let image_a = create_test_image(640, 480, Rgb([100, 110, 120]));
+        let image_b = create_test_image(420, 560, Rgb([10, 200, 90]));
+
+        let batched = processor
+            .preprocess(&[image_a.clone(), image_b.clone()], &config)
+            .unwrap();
+        let single_a = processor.preprocess(&[image_a], &config).unwrap();
+        let single_b = processor.preprocess(&[image_b], &config).unwrap();
+
+        // Feature token counts line up per image.
+        assert_eq!(
+            batched.feature_token_counts,
+            vec![
+                single_a.feature_token_counts[0],
+                single_b.feature_token_counts[0]
+            ]
+        );
+
+        // encoder_input rows: batch == [single_a rows ++ single_b rows].
+        let pa = single_a.encoder_input.shape()[0];
+        let pb = single_b.encoder_input.shape()[0];
+        assert_eq!(batched.encoder_input.shape()[0], pa + pb);
+        assert_eq!(
+            batched
+                .encoder_input
+                .slice_axis(Axis(0), Slice::from(0..pa))
+                .to_owned(),
+            single_a.encoder_input
+        );
+        assert_eq!(
+            batched
+                .encoder_input
+                .slice_axis(Axis(0), Slice::from(pa..pa + pb))
+                .to_owned(),
+            single_b.encoder_input
+        );
+
+        // image_grid_thw rows: batch == [single_a row ++ single_b row].
+        let grid = |inputs: &PreprocessedEncoderInputs| match inputs
+            .model_specific
+            .get("image_grid_thw")
+        {
+            Some(ModelSpecificValue::IntTensor { data, .. }) => data.clone(),
+            other => panic!("expected image_grid_thw IntTensor, got {other:?}"),
+        };
+        let mut expected_grid = grid(&single_a);
+        expected_grid.extend(grid(&single_b));
+        assert_eq!(grid(&batched), expected_grid);
+    }
+
     #[test]
     fn test_qwen3_vl_preprocess_video() {
         let processor = Qwen3VLProcessor::new();

crates/multimodal/src/vision/scratch.rs

@@ -0,0 +1,58 @@
+//! Recycling pools for the large per-image vision buffers.
+//!
+//! The preprocess pipeline allocates tens of MB per image (the [C, H, W] f32
+//! tensor and the batched patch buffer are each large).
+//! Freshly-allocated buffers of this size bypass the allocator's reuse paths
+//! (glibc caps non-main-arena chunks at 64 MB and mmaps anything larger or
+//! colder), so every image pays tens of thousands of minor page faults; the
+//! fault path serializes process-wide and caps the data plane's effective
+//! parallelism. Recycling keeps the pages mapped and hot.
+//!
+//! A lock-free thread-local pool serves same-thread take/give (preprocess
+//! internals run on blocking-pool threads). The pool is capped to bound
+//! residency; buffers beyond the cap are dropped.
+
+use std::cell::RefCell;
+
+/// Max recycled buffers kept per thread per class; excess is dropped. The
+/// vision path holds at most a couple of live tensors per request, so a small
+/// cap captures same-thread reuse.
+const MAX_THREAD_POOLED: usize = 2;
+
+thread_local! {
+    static F32_LOCAL: RefCell<Vec<Vec<f32>>> = const { RefCell::new(Vec::new()) };
+}
+
+macro_rules! pool_impl {
+    ($take_cap:ident, $give:ident, $ty:ty, $local:ident) => {
+        /// Take an empty `Vec` with at least `cap` capacity, reusing pooled storage.
+        pub fn $take_cap(cap: usize) -> Vec<$ty> {
+            let mut v = $local.with(|p| p.borrow_mut().pop()).unwrap_or_default();
+            v.clear();
+            v.reserve(cap);
+            v
+        }
+
+        /// Return a buffer for reuse by a later same-thread take.
+        pub fn $give(v: Vec<$ty>) {
+            if v.capacity() == 0 {
+                return;
+            }
+            $local.with(|p| {
+                let mut p = p.borrow_mut();
+                if p.len() < MAX_THREAD_POOLED {
+                    p.push(v);
+                }
+            });
+        }
+    };
+}
+
+pool_impl!(take_f32_cap, give_f32, f32, F32_LOCAL);
+
+/// Take a zero-filled `Vec<f32>` of exactly `len`, reusing pooled storage.
+pub fn take_f32(len: usize) -> Vec<f32> {
+    let mut v = take_f32_cap(len);
+    v.resize(len, 0.0);
+    v
+}

crates/multimodal/src/vision/transforms.rs

@@ -13,7 +13,10 @@ use image::{imageops::FilterType, DynamicImage, GenericImageView, Rgb, RgbImage}
 use ndarray::{s, Array3, Array4};
 use thiserror::Error;
 
-use super::execution::{scope as parallel_scope, task_count};
+use super::{
+    execution::{scope as parallel_scope, task_count},
+    scratch,
+};
 
 /// Errors that can occur during image transformations.
 #[derive(Error, Debug)]
@@ -151,7 +154,8 @@ fn build_planar_tensor(
     bias: [f32; 3],
 ) -> Array3<f32> {
     let pixels = h * w;
-    let mut data = vec![0.0f32; 3 * pixels];
+    // Pooled: this large per-image buffer is the data plane's hottest allocation.
+    let mut data = scratch::take_f32(3 * pixels);
     let (r_plane, rest) = data.split_at_mut(pixels);
     let (g_plane, b_plane) = rest.split_at_mut(pixels);
 

crates/protocols/src/worker.rs

@@ -723,6 +723,9 @@ impl WorkerSpec {
 /// - `Shm`: use same-host `/dev/shm` when SMG can write it.
 /// - `Auto`: use `/dev/shm` only when the receiving worker is verified to share
 ///   SMG's `/dev/shm`; otherwise fall back to inline.
+/// - `Rdma`: use the NIXL RDMA pixel lane for large tensors (requires the
+///   `mm-rdma` build feature + NIXL); falls back to inline when RDMA is
+///   unavailable.
 #[derive(
     Debug, Clone, Copy, PartialEq, Eq, Hash, Default, Serialize, Deserialize, schemars::JsonSchema,
 )]
@@ -732,15 +735,17 @@ pub enum TransportMode {
     Inline,
     Shm,
     Auto,
+    Rdma,
 }
 
 impl TransportMode {
-    /// Parse from a case-insensitive string (`inline` | `shm` | `auto`).
+    /// Parse from a case-insensitive string (`inline` | `shm` | `auto` | `rdma`).
     pub fn parse(value: &str) -> Option<Self> {
         match value.trim().to_ascii_lowercase().as_str() {
             "inline" => Some(Self::Inline),
             "shm" => Some(Self::Shm),
             "auto" => Some(Self::Auto),
+            "rdma" => Some(Self::Rdma),
             _ => None,
         }
     }
@@ -751,6 +756,7 @@ impl TransportMode {
             Self::Inline => "inline",
             Self::Shm => "shm",
             Self::Auto => "auto",
+            Self::Rdma => "rdma",
         }
     }
 }
@@ -766,7 +772,7 @@ impl std::str::FromStr for TransportMode {
 
     fn from_str(s: &str) -> Result<Self, Self::Err> {
         Self::parse(s)
-            .ok_or_else(|| format!("invalid transport mode '{s}'; expected inline|shm|auto"))
+            .ok_or_else(|| format!("invalid transport mode '{s}'; expected inline|shm|auto|rdma"))
     }
 }
 

docs/reference/configuration.md

@@ -949,7 +949,7 @@ router config / CLI flag → `SMG_MM_*` environment variable → built-in defaul
 
 | CLI Flag | Config / WorkerSpec field | Default | Description |
 |----------|---------------------------|---------|-------------|
-| `--multimodal-tensor-transport` | `multimodal_tensor_transport` | `inline` | Transport for large MM tensors: `inline` (gRPC bytes), `shm` (use `/dev/shm` whenever the router can write it — the operator asserts co-location), 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. |
+| `--multimodal-tensor-transport` | `multimodal_tensor_transport` | `inline` | Transport for large MM tensors: `inline` (gRPC bytes), `shm` (use `/dev/shm` whenever the router can write it — the operator asserts co-location), `auto` (use `/dev/shm` only when the worker is *verified* to share it), or `rdma` (pull pixels over the NIXL RDMA lane; requires the `mm-rdma` build feature + NIXL, and falls back to inline when unavailable). 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. |
 | `--multimodal-shm-min-bytes` | `multimodal_shm_min_bytes` | `65536` | Minimum tensor size (bytes) before the SHM path is used; smaller tensors stay inline. |
 
 Both settings can be overridden per worker via the matching `WorkerSpec` fields