From d6d7940f7056767917083a90d01b19b442c44d32 Mon Sep 17 00:00:00 2001 From: chenht2022 Date: Sun, 28 Jun 2026 13:49:37 +0000 Subject: [PATCH 1/4] feat(multimodal): optimize EPD encode routing Optimize multimodal tensor handling for EPD encode routing. Add optional fast-path transfer support, scratch buffer reuse, and model preprocessing/serialization improvements. Signed-off-by: chenht2022 --- crates/multimodal/src/vision/mod.rs | 1 + .../src/vision/processors/kimi_k25.rs | 36 +- .../src/vision/processors/qwen3_vl.rs | 68 ++ crates/multimodal/src/vision/scratch.rs | 58 ++ crates/multimodal/src/vision/transforms.rs | 3 +- .../tokenspeed/encoder_servicer.py | 40 +- .../tokenspeed/rdma_pixel.py | 339 +++++++++ .../smg_grpc_servicer/tokenspeed/servicer.py | 18 +- model_gateway/Cargo.toml | 5 + model_gateway/src/observability/metrics.rs | 2 +- model_gateway/src/routers/grpc/client.rs | 8 +- model_gateway/src/routers/grpc/epd_encode.rs | 3 +- model_gateway/src/routers/grpc/mm_rdma/mod.rs | 16 + .../src/routers/grpc/mm_rdma/nixl.rs | 698 ++++++++++++++++++ .../src/routers/grpc/mm_rdma/stub.rs | 9 + model_gateway/src/routers/grpc/mod.rs | 1 + .../src/routers/grpc/multimodal/config.rs | 5 + .../src/routers/grpc/multimodal/mod.rs | 1 + .../routers/grpc/multimodal/pixel_cache.rs | 256 +++++++ .../src/routers/grpc/multimodal/process.rs | 205 +++-- .../src/routers/grpc/multimodal/serialize.rs | 3 + .../src/routers/grpc/proto_wrapper.rs | 114 ++- 22 files changed, 1792 insertions(+), 97 deletions(-) create mode 100644 crates/multimodal/src/vision/scratch.rs create mode 100644 grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py create mode 100644 model_gateway/src/routers/grpc/mm_rdma/mod.rs create mode 100644 model_gateway/src/routers/grpc/mm_rdma/nixl.rs create mode 100644 model_gateway/src/routers/grpc/mm_rdma/stub.rs create mode 100644 model_gateway/src/routers/grpc/multimodal/pixel_cache.rs diff --git a/crates/multimodal/src/vision/mod.rs b/crates/multimodal/src/vision/mod.rs index 0afea824fb..83143818ba 100644 --- a/crates/multimodal/src/vision/mod.rs +++ b/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 diff --git a/crates/multimodal/src/vision/processors/kimi_k25.rs b/crates/multimodal/src/vision/processors/kimi_k25.rs index 16f6798515..a2ea599a9e 100644 --- a/crates/multimodal/src/vision/processors/kimi_k25.rs +++ b/crates/multimodal/src/vision/processors/kimi_k25.rs @@ -156,7 +156,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 = crate::vision::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 +195,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, patch_size: usize) -> Vec { + /// 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, patch_size: usize, out: &mut Vec) { 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 +223,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 +253,18 @@ impl VisionPreProcessor for KimiK25Processor { let mean = config.get_image_mean(); let std = config.get_image_std(); - let mut all_patches: Vec = 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 = crate::vision::scratch::take_f32_cap(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()); @@ -280,8 +289,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(); + crate::vision::scratch::give_f32(storage); patches_per_image.push(num_patches as i64); } diff --git a/crates/multimodal/src/vision/processors/qwen3_vl.rs b/crates/multimodal/src/vision/processors/qwen3_vl.rs index e163967b47..05bda801a5 100644 --- a/crates/multimodal/src/vision/processors/qwen3_vl.rs +++ b/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(); diff --git a/crates/multimodal/src/vision/scratch.rs b/crates/multimodal/src/vision/scratch.rs new file mode 100644 index 0000000000..bf22234c89 --- /dev/null +++ b/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>> = 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` of exactly `len`, reusing pooled storage. +pub fn take_f32(len: usize) -> Vec { + let mut v = take_f32_cap(len); + v.resize(len, 0.0); + v +} diff --git a/crates/multimodal/src/vision/transforms.rs b/crates/multimodal/src/vision/transforms.rs index fde6e1a2a3..875773136c 100644 --- a/crates/multimodal/src/vision/transforms.rs +++ b/crates/multimodal/src/vision/transforms.rs @@ -151,7 +151,8 @@ fn build_planar_tensor( bias: [f32; 3], ) -> Array3 { 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 = crate::vision::scratch::take_f32(3 * pixels); let (r_plane, rest) = data.split_at_mut(pixels); let (g_plane, b_plane) = rest.split_at_mut(pixels); diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py index 759bef3a19..29019981c0 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py @@ -22,6 +22,7 @@ tokenspeed_scheduler_pb2, ) +from smg_grpc_servicer.tokenspeed.rdma_pixel import RdmaPixelPuller from smg_grpc_servicer.tokenspeed.servicer import TokenSpeedSchedulerServicer if TYPE_CHECKING: @@ -79,6 +80,11 @@ def __init__( # Spatial merge factor for post-merge token counts (Qwen vision default 2). self._merge_size = self._resolve_merge_size() + self._rdma_pixel_puller = RdmaPixelPuller( + agent_name=f"smg-encode-{self._bootstrap_host}-{self._bootstrap_port}", + log_prefix="EPD RDMA", + ) + self.async_llm.auto_create_handle_loop() logger.info("TokenSpeedEncoderServicer initialized") @@ -87,7 +93,7 @@ def _resolve_merge_size(self) -> int: vision_config = getattr(hf_config, "vision_config", None) return int(getattr(vision_config, "spatial_merge_size", 2) or 2) - def _items_from_proto(self, mm_inputs): + def _items_from_proto(self, mm_inputs, bootstrap_room: int = 0): """Reconstruct the engine MultimodalDataItem(s) for the encode worker. Unlike the prefill leg, the encode worker NEEDS each item's encoder_input @@ -111,14 +117,28 @@ def _items_from_proto(self, mm_inputs): # the 19-77MB pixels. The content hash is computed on the real bytes # before the swap and pre-set on the item. td = item_proto.encoder_input - feature = TokenSpeedSchedulerServicer._tensor_from_proto(td, cast_to=model_dtype) - feat_hash = None - if self._pixel_shm: - from tokenspeed.runtime.multimodal.hash import hash_feature - from tokenspeed.runtime.multimodal.shm_transport import ShmTensorHandle - - feat_hash = hash_feature(feature) - feature = ShmTensorHandle.publish(feature) + if td.WhichOneof("payload") == "remote": + # EPD RDMA: pull pixels from the gateway's exported NIXL memory. + # With EPD_PIXEL_SHM, the received slot is published directly to + # scheduler SHM so the scheduler ingest path still avoids pickle + # copies of the full pixel tensor. + feature, feat_hash = self._rdma_pixel_puller.feature_from_remote( + td, + explicit_room=bootstrap_room, + cast_to=model_dtype, + publish_shm=self._pixel_shm, + ) + else: + feature = TokenSpeedSchedulerServicer._tensor_from_proto(td, cast_to=model_dtype) + feat_hash = None + if self._pixel_shm: + from tokenspeed.runtime.multimodal.hash import hash_feature + from tokenspeed.runtime.multimodal.shm_transport import ( + ShmTensorHandle, + ) + + feat_hash = hash_feature(feature) + feature = ShmTensorHandle.publish(feature) model_specific = { name: TokenSpeedSchedulerServicer._tensor_from_proto(t, cast_to=model_dtype) for name, t in item_proto.model_specific_tensors.items() @@ -214,7 +234,7 @@ async def Encode(self, request, context): def _build_encode_request(self, request, bootstrap_room): """Proto -> engine EncodeRequest (the expensive per-image parse).""" - items = self._items_from_proto(request.mm_inputs) + items = self._items_from_proto(request.mm_inputs, bootstrap_room) EncodeRequest = _lazy_encode_request() return EncodeRequest( diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py new file mode 100644 index 0000000000..2782cc9b80 --- /dev/null +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py @@ -0,0 +1,339 @@ +"""Shared NIXL pixel-payload puller for TokenSpeed multimodal inputs.""" + +from __future__ import annotations + +import logging +import os +import queue +import socket +import threading +import time +from dataclasses import dataclass + +logger = logging.getLogger(__name__) + +_DESCRIPTOR_MAGIC = b"SMGRDMA1" +_GEN_BYTES = 8 +_FRAME_BYTES = 2 * _GEN_BYTES +_GATEWAY_AGENT_NAME = "smg-gateway-encode" +_LOCAL_IP_CACHE: dict[str, bool] = {} + + +@dataclass(frozen=True) +class _RemotePixelDescriptor: + remote_addr: int + expected_gen: int + room: int + port: int + ip: str + + +def _ip_is_local(ip: str) -> bool: + """True iff ``ip`` is assigned to a local interface of this host.""" + cached = _LOCAL_IP_CACHE.get(ip) + if cached is not None: + return cached + result = False + if ip in ("localhost", "::1") or ip.startswith("127."): + result = True + else: + for fam in (socket.AF_INET, socket.AF_INET6): + try: + sock = socket.socket(fam, socket.SOCK_DGRAM) + except OSError: + continue + try: + sock.bind((ip, 0)) + result = True + except OSError: + pass + finally: + sock.close() + if result: + break + _LOCAL_IP_CACHE[ip] = result + return result + + +def _parse_descriptor(td, explicit_room: int | None) -> _RemotePixelDescriptor: + """Parse both the current room-carrying descriptor and the legacy EPD form.""" + desc = bytes(td.remote.descriptor) + if desc.startswith(_DESCRIPTOR_MAGIC): + min_len = len(_DESCRIPTOR_MAGIC) + 8 + _GEN_BYTES + 8 + 2 + if len(desc) < min_len: + raise ValueError("remote descriptor too short") + off = len(_DESCRIPTOR_MAGIC) + remote_addr = int.from_bytes(desc[off : off + 8], "little") + off += 8 + expected_gen = int.from_bytes(desc[off : off + _GEN_BYTES], "little") + off += _GEN_BYTES + room = int.from_bytes(desc[off : off + 8], "little", signed=True) + off += 8 + port = int.from_bytes(desc[off : off + 2], "little") + off += 2 + ip = desc[off:].decode() + if explicit_room is not None and int(explicit_room) != room: + raise ValueError( + f"remote descriptor room mismatch: descriptor={room} request={int(explicit_room)}" + ) + return _RemotePixelDescriptor( + remote_addr=remote_addr, + expected_gen=expected_gen, + room=room, + port=port, + ip=ip, + ) + + if explicit_room is None: + raise ValueError( + "legacy remote descriptor lacks room; gateway must emit SMGRDMA1 descriptor" + ) + if len(desc) < 8 + _GEN_BYTES + 2: + raise ValueError("remote descriptor too short") + return _RemotePixelDescriptor( + remote_addr=int.from_bytes(desc[:8], "little"), + expected_gen=int.from_bytes(desc[8 : 8 + _GEN_BYTES], "little"), + room=int(explicit_room), + port=int.from_bytes(desc[8 + _GEN_BYTES : 10 + _GEN_BYTES], "little"), + ip=desc[10 + _GEN_BYTES :].decode(), + ) + + +class RdmaPixelPuller: + """Persistent NIXL READ agent for gateway-exported multimodal pixel buffers.""" + + def __init__(self, *, agent_name: str, log_prefix: str): + self._log_prefix = log_prefix + self._nixl_agent = None + self._rdma_md_ready = set() + self._rdma_md_lock = threading.Lock() + self._landing = None + self._landing_np = None + self._landing_base = 0 + self._landing_slot_bytes = 0 + self._landing_free = None + + if os.environ.get("SMG_MM_PIXEL_RDMA") not in ("1", "true"): + return + + try: + try: + from nixl_cu13._api import nixl_agent, nixl_agent_config + except ImportError: + from nixl._api import nixl_agent, nixl_agent_config + import torch + + self._nixl_agent = nixl_agent( + agent_name, + nixl_agent_config( + enable_listen_thread=True, + listen_port=0, + backends=["UCX"], + ), + ) + + slot_bytes = int(os.environ.get("SMG_RDMA_SLOT_BYTES", 32 * 1024 * 1024)) + n_slots = int(os.environ.get("SMG_RDMA_LANDING_SLOTS", 64)) + self._landing = torch.empty( + n_slots * slot_bytes, + dtype=torch.uint8, + pin_memory=True, + ) + self._landing_np = self._landing.numpy() + self._landing_base = self._landing.data_ptr() + self._landing_slot_bytes = slot_bytes + reg = self._nixl_agent.get_reg_descs( + [(self._landing_base, n_slots * slot_bytes, 0, "")], + "DRAM", + ) + self._nixl_agent.register_memory(reg) + self._landing_free = queue.Queue() + for i in range(n_slots): + self._landing_free.put(i) + logger.info( + "%s: puller up (landing %d slots x %d B)", + self._log_prefix, + n_slots, + slot_bytes, + ) + except Exception as exc: # noqa: BLE001 + logger.error("%s: NIXL init failed (%s); inline only", self._log_prefix, exc) + self._nixl_agent = None + self._landing_free = None + + def _ensure_remote_ready(self, ip: str, port: int, remote, room: int) -> None: + """One-time metadata handshake per gateway listener.""" + key = (ip, port) + if key in self._rdma_md_ready: + return + with self._rdma_md_lock: + if key in self._rdma_md_ready: + return + + agent = self._nixl_agent + agent.fetch_remote_metadata(_GATEWAY_AGENT_NAME, ip, port) + + send_md_env = os.environ.get("SMG_RDMA_SEND_MD") + if send_md_env in ("1", "true"): + do_send_md = True + elif send_md_env in ("0", "false"): + do_send_md = False + else: + do_send_md = _ip_is_local(ip) + logger.info( + "%s: gateway %s:%s local=%s send_md=%s (env=%s) -- one-time md handshake", + self._log_prefix, + ip, + port, + _ip_is_local(ip), + do_send_md, + send_md_env, + ) + if do_send_md: + agent.send_local_metadata(ip, port) + + ready = False + for _ in range(5000): + if agent.check_remote_metadata(_GATEWAY_AGENT_NAME, remote): + ready = True + break + time.sleep(0.001) + if not ready: + raise RuntimeError(f"NIXL remote metadata not ready room={room}") + self._rdma_md_ready.add(key) + + def feature_from_remote( + self, + td, + *, + explicit_room: int | None, + cast_to, + publish_shm: bool = False, + ): + """Read a remote TensorData payload and return ``(feature, content_hash)``.""" + import numpy as np + import torch + + transport = getattr(td.remote, "transport", "") + if transport and transport != "nixl": + raise ValueError(f"unsupported remote tensor transport: {transport!r}") + + agent = self._nixl_agent + if agent is None or self._landing_free is None: + raise RuntimeError(f"{self._log_prefix}: remote payload but landing pool unavailable") + + descriptor = _parse_descriptor(td, explicit_room) + nbytes = int(td.remote.nbytes) + if nbytes + _FRAME_BYTES > self._landing_slot_bytes: + raise ValueError( + f"remote pixel {nbytes}B (+{_FRAME_BYTES}B gen frame) exceeds " + f"landing slot {self._landing_slot_bytes}B" + ) + + remote = agent.get_xfer_descs( + [(descriptor.remote_addr, nbytes + _FRAME_BYTES, 0)], + "DRAM", + ) + self._ensure_remote_ready( + descriptor.ip, + descriptor.port, + remote, + descriptor.room, + ) + + wait_budget = float(os.environ.get("SMG_RDMA_LANDING_WAIT_S", 120)) + t_acq = time.monotonic() + slot = None + while slot is None: + try: + slot = self._landing_free.get(timeout=5) + except queue.Empty: + waited = time.monotonic() - t_acq + if waited >= wait_budget: + raise RuntimeError( + f"{self._log_prefix}: landing-ring starvation for {waited:.0f}s " + f"(room={descriptor.room}); raise SMG_RDMA_LANDING_SLOTS / " + f"SMG_RDMA_LANDING_WAIT_S" + ) from None + logger.warning( + "%s: landing ring exhausted for %.0fs (room=%s, qsize=%d); backpressuring", + self._log_prefix, + waited, + descriptor.room, + self._landing_free.qsize(), + ) + + try: + off = slot * self._landing_slot_bytes + laddr = self._landing_base + off + local = agent.get_xfer_descs([(laddr, nbytes + _FRAME_BYTES, 0)], "DRAM") + handle = agent.initialize_xfer( + "READ", + local, + remote, + _GATEWAY_AGENT_NAME, + str(descriptor.room).encode(), + ) + read_deadline = time.monotonic() + float(os.environ.get("SMG_RDMA_READ_TIMEOUT_S", 60)) + spins = 0 + state = agent.transfer(handle) + while state in ("PROC", "IN_PROG"): + spins += 1 + if spins > 64: + time.sleep(0.0005) + if time.monotonic() > read_deadline: + agent.release_xfer_handle(handle) + raise RuntimeError(f"NIXL READ timed out room={descriptor.room}") + state = agent.check_xfer_state(handle) + agent.release_xfer_handle(handle) + if state != "DONE": + raise RuntimeError(f"NIXL READ state={state} room={descriptor.room}") + + shape = list(td.shape) + itemsize = 2 if td.dtype == "bfloat16" else np.dtype(td.dtype).itemsize + expected = itemsize + for dim in shape: + expected *= dim + if nbytes != expected: + raise ValueError( + f"remote pixel size mismatch: nbytes={nbytes} expected={expected} " + f"(shape={shape} dtype={td.dtype})" + ) + + hdr = int.from_bytes( + self._landing_np[off : off + _GEN_BYTES].tobytes(), + "little", + ) + trl = int.from_bytes( + self._landing_np[off + _GEN_BYTES + nbytes : off + _FRAME_BYTES + nbytes].tobytes(), + "little", + ) + if hdr != descriptor.expected_gen or trl != descriptor.expected_gen: + raise ValueError( + f"{self._log_prefix}: slot generation mismatch room={descriptor.room} " + f"expected={descriptor.expected_gen} header={hdr} trailer={trl} " + f"(slot recycled/torn under a live READ -> wrong-image guard)" + ) + + slot_np = self._landing_np[off + _GEN_BYTES : off + _GEN_BYTES + nbytes] + if td.dtype == "bfloat16": + tensor = torch.from_numpy(slot_np.view(np.uint16).reshape(shape)).view( + torch.bfloat16 + ) + else: + tensor = torch.from_numpy(slot_np.view(np.dtype(td.dtype)).reshape(shape)) + copied = False + if cast_to is not None and tensor.dtype != cast_to and tensor.is_floating_point(): + tensor = tensor.to(cast_to) + copied = True + + if publish_shm: + from tokenspeed.runtime.multimodal.hash import hash_feature + from tokenspeed.runtime.multimodal.shm_transport import ShmTensorHandle + + feat_hash = hash_feature(tensor) + return ShmTensorHandle.publish(tensor), feat_hash + + return (tensor if copied else tensor.clone()), None + finally: + self._landing_free.put(slot) diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py index 968934940e..a32ab06672 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py @@ -41,6 +41,7 @@ from smg_grpc_servicer.kv_events import endpoint_for_rank, stream_kv_events from smg_grpc_servicer.tokenspeed.health_servicer import TokenSpeedHealthServicer from smg_grpc_servicer.tokenspeed.kv_events import resolve_kv_events_config +from smg_grpc_servicer.tokenspeed.rdma_pixel import RdmaPixelPuller if TYPE_CHECKING: # Type-only — keeps these out of the cold-path graph when the servicer is @@ -131,6 +132,13 @@ def __init__( # Resolved ZMQ KV-events endpoint, or None when the worker was not # launched with --kv-events-config (SubscribeKvEvents → UNIMPLEMENTED). self._kv_events_config = resolve_kv_events_config(server_args) + self._rdma_pixel_puller = RdmaPixelPuller( + agent_name=( + f"smg-scheduler-{getattr(server_args, 'host', 'unknown')}-" + f"{getattr(server_args, 'port', 'unknown')}" + ), + log_prefix="TokenSpeed RDMA", + ) # Drive AsyncLLM's output-dispatch loop. This is idempotent — the # first caller creates the handle loop; subsequent callers (including @@ -1027,7 +1035,15 @@ def _mm_inputs_from_itemized_proto( if item_proto.HasField("encoder_input"): feature_started = time.perf_counter() if LOG_MM_TIMING else None encoder_input = item_proto.encoder_input - feature = self._feature_from_proto(encoder_input, cast_to=model_dtype) + if encoder_input.WhichOneof("payload") == "remote": + feature, _ = self._rdma_pixel_puller.feature_from_remote( + encoder_input, + explicit_room=None, + cast_to=model_dtype, + publish_shm=False, + ) + else: + feature = self._feature_from_proto(encoder_input, cast_to=model_dtype) feature_elapsed_ms = ( (time.perf_counter() - feature_started) * 1000 if feature_started is not None diff --git a/model_gateway/Cargo.toml b/model_gateway/Cargo.toml index b075ca3bb7..96e0a163f3 100644 --- a/model_gateway/Cargo.toml +++ b/model_gateway/Cargo.toml @@ -19,6 +19,7 @@ default = ["grpc-client"] grpc-client = [] grpc-server = [] opencv-video = ["llm-multimodal/opencv-video"] +mm-rdma = ["dep:nixl-sys"] vendored-openssl = ["openssl/vendored"] @@ -35,6 +36,9 @@ name = "amg" path = "src/main.rs" [dependencies] +# Multimodal RDMA pixel transport. Optional because nixl-sys runs bindgen at +# build time; ordinary gateway builds should not require a NIXL/clang setup. +nixl-sys = { version = "0.10.1", features = ["stub-api"], optional = true } # Workspace dependencies (simple) anyhow.workspace = true async-trait.workspace = true @@ -43,6 +47,7 @@ chrono.workspace = true dashmap.workspace = true futures.workspace = true http.workspace = true +lru.workspace = true parking_lot.workspace = true prost.workspace = true prost-types.workspace = true diff --git a/model_gateway/src/observability/metrics.rs b/model_gateway/src/observability/metrics.rs index 711d999d48..18595fe19d 100644 --- a/model_gateway/src/observability/metrics.rs +++ b/model_gateway/src/observability/metrics.rs @@ -647,7 +647,7 @@ impl Metrics { .increment(1); } - /// Record one multimodal tensor sent over `path` ("inline"|"shm") for `runtime`. + /// Record one multimodal tensor sent over `path` ("inline"|"shm"|"remote") for `runtime`. pub fn record_mm_tensor(runtime: &'static str, path: &'static str, nbytes: usize) { counter!("smg_mm_tensors_total", "runtime" => runtime, "path" => path).increment(1); counter!("smg_mm_tensor_bytes_total", "runtime" => runtime, "path" => path) diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index 0c3e9dfc06..990a291d11 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -534,7 +534,9 @@ impl GrpcClient { } Self::TokenSpeed(client) => { let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { - MultimodalData::TokenSpeed(data) => data.into_proto(), + MultimodalData::TokenSpeed(data) => { + data.try_export_encoder_inputs_nixl_remote().into_proto() + } _ => unreachable!("caller guarantees matching variant"), }); finish_tokenspeed_request(tokenspeed_mm, |mm| { @@ -626,7 +628,9 @@ impl GrpcClient { } Self::TokenSpeed(client) => { let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { - MultimodalData::TokenSpeed(data) => data.into_proto(), + MultimodalData::TokenSpeed(data) => { + data.try_export_encoder_inputs_nixl_remote().into_proto() + } _ => unreachable!("caller guarantees matching variant"), }); finish_tokenspeed_request(tokenspeed_mm, |mm| { diff --git a/model_gateway/src/routers/grpc/epd_encode.rs b/model_gateway/src/routers/grpc/epd_encode.rs index edae614e00..b37f1d7084 100644 --- a/model_gateway/src/routers/grpc/epd_encode.rs +++ b/model_gateway/src/routers/grpc/epd_encode.rs @@ -104,10 +104,11 @@ impl PreparedEncodeItem { shm_min_bytes, cleanup_on_drop, } => { - let item = item + let mut item = item .take() .ok_or_else(|| "encode item was already dispatched".to_string())?; *cleanup_on_drop = false; + item.encoder_input = item.encoder_input.try_export_nixl_remote(bootstrap_room); let request = tokenspeed_encoder::EncodeRequest { request_id: format!("encode-{}", Uuid::now_v7()), mm_inputs: Some( diff --git a/model_gateway/src/routers/grpc/mm_rdma/mod.rs b/model_gateway/src/routers/grpc/mm_rdma/mod.rs new file mode 100644 index 0000000000..4555bb74f9 --- /dev/null +++ b/model_gateway/src/routers/grpc/mm_rdma/mod.rs @@ -0,0 +1,16 @@ +//! Optional multimodal pixel RDMA transport. +//! +//! The default build uses a no-op shim so ordinary gateway builds do not need +//! NIXL headers/bindgen. Enable the `mm-rdma` Cargo feature to compile the NIXL +//! implementation. + +#[cfg(feature = "mm-rdma")] +mod nixl; + +#[cfg(not(feature = "mm-rdma"))] +mod stub; + +#[cfg(feature = "mm-rdma")] +pub(crate) use nixl::*; +#[cfg(not(feature = "mm-rdma"))] +pub(crate) use stub::*; diff --git a/model_gateway/src/routers/grpc/mm_rdma/nixl.rs b/model_gateway/src/routers/grpc/mm_rdma/nixl.rs new file mode 100644 index 0000000000..f708b6a433 --- /dev/null +++ b/model_gateway/src/routers/grpc/mm_rdma/nixl.rs @@ -0,0 +1,698 @@ +//! EPD RDMA pixel transport: the gateway exports each image's serialized pixel +//! buffer over NIXL (UCX/RoCE) so the encode worker PULLs it (one-sided READ), +//! instead of shipping the large pixel buffer inline in the Encode gRPC frame. +//! On the inline path, serializing it burns gateway CPU and starves the encode +//! worker's GPU -- the EPD bottleneck this transport removes. +//! +//! Gated behind `SMG_MM_PIXEL_RDMA`; every NIXL call is reached only when the gate +//! is on (a missing `libnixl_capi.so` aborts on the first stub call, not a no-op), +//! and any export failure falls back to the inline payload so EPD never hard-fails. +//! +//! v2 (pre-registered pool): one host-DRAM arena is registered with NIXL ONCE at +//! init and sub-divided into fixed slots. Per image the hot path only LEASES a free +//! slot, frames the pixels as `[gen u64][pixels][gen u64]`, and ships +//! `[magic 8B][slot_addr u64 LE][gen u64 LE][room i64 LE][port u16 LE][ip utf8]` +//! -- NO per-image register_memory and NO growing agent metadata. The worker fetches +//! the gateway's (now fixed) metadata ONCE, then READs slot offsets into its own +//! pre-registered landing pool. The transfer notif is tagged with the bootstrap_room; +//! this reaper consumes it to return the slot to the free list (a TTL sweep is the +//! lost-notif net). This removes the v1 per-image control+registration overhead that +//! made the RDMA path slower than inline. + +use std::{ + sync::{ + atomic::{AtomicU64, Ordering}, + OnceLock, + }, + time::{Duration, Instant}, +}; + +use dashmap::DashMap; +use nixl_sys::{ + Agent, AgentConfig, Backend, MemType, MemoryRegion, NixlDescriptor, NotificationMap, OptArgs, + RegistrationHandle, +}; +use parking_lot::Mutex; +use tracing::{debug, error}; + +/// `SMG_MM_PIXEL_RDMA` gate (cached in a process-wide `OnceLock`). +pub(crate) fn rdma_enabled() -> bool { + static ENABLED: OnceLock = OnceLock::new(); + *ENABLED.get_or_init(|| { + matches!( + std::env::var("SMG_MM_PIXEL_RDMA").as_deref(), + Ok("1") | Ok("true") + ) + }) +} + +/// Reaper poll interval. +const REAPER_TICK: Duration = Duration::from_millis(20); +/// Default arena geometry: 64 slots x 32 MiB = 2 GiB host DRAM. A slot must hold one +/// image's bf16 pixel buffer; raise SLOT_BYTES if larger images overflow it. +const DEFAULT_POOL_SLOTS: usize = 64; +const DEFAULT_SLOT_BYTES: usize = 32 * 1024 * 1024; + +/// Per-lease generation framing. Each leased slot is written as +/// `[gen u64 LE][payload][gen u64 LE]` and the descriptor carries the same `gen`, so +/// the worker can detect a slot it READ after the gateway recycled and reused it (the +/// residual the TTL widening still leaves if the TTL is ever misconfigured / the two +/// sides' env drift): a recycled slot carries a strictly newer gen than the worker's +/// descriptor, and a READ torn by a concurrent reuse sees header != trailer. Both are +/// caught by requiring `header == trailer == descriptor.gen`, which makes correctness +/// independent of the TTL value (TTL stays purely a capacity knob). Worker-side parse +/// mirrors this in grpc_servicer encoder_servicer.py (`GEN`/`FRAME`). +const GEN_BYTES: usize = 8; +/// Header + trailer generation stamps bracketing the payload. +const FRAME_OVERHEAD: usize = 2 * GEN_BYTES; +/// Descriptor prefix for the version that carries `bootstrap_room` inline. +const DESCRIPTOR_MAGIC: &[u8; 8] = b"SMGRDMA1"; + +/// Read a seconds-valued env knob, falling back to `default`. +fn env_secs(key: &str, default: u64) -> u64 { + std::env::var(key) + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(default) +} + +/// The worst-case wall time the encode worker may hold a shipped descriptor before +/// and during its one-sided READ: it waits up to `SMG_RDMA_LANDING_WAIT_S` for a free +/// landing slot, then READs for up to `SMG_RDMA_READ_TIMEOUT_S`. These mirror the +/// encode servicer's own knobs (same env names, same defaults) so the two sides +/// cannot drift. The gateway must not reclaim a slot inside this window. +fn worker_max_hold() -> Duration { + Duration::from_secs( + env_secs("SMG_RDMA_LANDING_WAIT_S", 120) + env_secs("SMG_RDMA_READ_TIMEOUT_S", 60), + ) +} + +/// Fixed slack added to the derived `worker_max_hold()` when deriving the slot +/// TTL. A const rather than an env knob: it would only be a second tuning layer on +/// an already-derived value, it only ever widens the lost-notif leak window (a +/// capacity nit, never correctness -- the per-lease gen-framing in `GEN_BYTES` makes +/// a recycled-under-read slot detectable independent of the TTL), and 30s dwarfs any +/// Encode-RPC delivery jitter. `SMG_RDMA_SLOT_TTL_S` remains the explicit full-TTL +/// override. +const SLOT_TTL_SLACK: Duration = Duration::from_secs(30); + +/// How long a leased slot may live without a free-notif before the reaper +/// force-reclaims it (lost notif / dead worker). +/// +/// CORRECTNESS: this MUST exceed `worker_max_hold()` (plus slack for the Encode-RPC +/// delivery the gateway can't see), otherwise the TTL races the worker's still-valid +/// READ: the reaper returns the slot, the next image re-leases the SAME address, and +/// the late READ silently returns the WRONG image's pixels (no error, no NaN). Derived +/// by default (= worker_max_hold + the fixed `SLOT_TTL_SLACK`); `SMG_RDMA_SLOT_TTL_S` +/// overrides explicitly. The cost of the wider TTL is only that a genuinely lost notif +/// leaks one slot for longer -- a capacity nit, never a correctness one. +fn slot_ttl() -> Duration { + static TTL: OnceLock = OnceLock::new(); + *TTL.get_or_init(|| { + if let Ok(v) = std::env::var("SMG_RDMA_SLOT_TTL_S") { + if let Ok(secs) = v.parse::() { + return Duration::from_secs(secs); + } + } + worker_max_hold() + SLOT_TTL_SLACK + }) +} + +/// NIXL descriptor over the pre-registered arena (host DRAM). Registered once; the +/// backing allocation is leaked (process-lifetime) so `base` is stable forever. +#[derive(Debug)] +struct ArenaRegion { + base: usize, + size: usize, +} + +impl MemoryRegion for ArenaRegion { + // The NIXL trait declares this `unsafe`; the body just reinterprets a usize the + // crate denies unsafe_code globally, so allow it for this trait impl. + #[allow(unsafe_code)] + unsafe fn as_ptr(&self) -> *const u8 { + self.base as *const u8 + } + fn size(&self) -> usize { + self.size + } +} + +impl NixlDescriptor for ArenaRegion { + fn mem_type(&self) -> MemType { + MemType::Dram + } + fn device_id(&self) -> u64 { + 0 + } +} + +/// A leased slot, kept until the worker's READ-notif (tagged with bootstrap_room) +/// or the TTL sweep returns it to the free list. +struct OccSlot { + slot: u32, + at: Instant, +} + +/// Pure slot bookkeeping over the arena's raw memory: the free-list, the leased +/// (room -> slot) map, and the memcpy-into-slot. Deliberately holds NO NIXL handle so +/// the lease / reclaim / reuse policy (where the TTL-vs-READ race lives) is unit- +/// testable without a registered agent or any RDMA hardware. +struct SlotPool { + /// Raw base address of the leaked arena allocation (usize => Send+Sync). + base: usize, + slot_bytes: usize, + n_slots: usize, + /// Available slot indices. + free: Mutex>, + /// Leased slots keyed by bootstrap_room (free-on-notif + TTL reclaim). + occupied: DashMap, + /// Monotonic per-lease generation stamp (see GEN_BYTES). Starts at 1 so the + /// zero-initialized arena (a never-written slot reads gen 0) can never match a + /// real descriptor's gen. + gen: AtomicU64, +} + +impl SlotPool { + fn new(base: usize, slot_bytes: usize, n_slots: usize) -> SlotPool { + SlotPool { + base, + slot_bytes, + n_slots, + free: Mutex::new((0..n_slots as u32).collect()), + occupied: DashMap::new(), + gen: AtomicU64::new(1), + } + } + + /// Lease a free slot for `room`, frame `bytes` as `[gen][payload][gen]` into it, and + /// record the lease at `now`. Returns `(slot, slot_addr, gen)`, or `None` if the + /// framed image exceeds `slot_bytes` or no slot is free (caller -> inline). Recording + /// the lease here (rather than in the caller) keeps the slot and its TTL timestamp + /// atomic w.r.t. the write. `gen` is shipped in the descriptor and re-checked by the + /// worker against the slot's stamps to reject a recycled/torn READ. + fn lease_and_write(&self, room: i64, bytes: &[u8], now: Instant) -> Option<(u32, u64, u64)> { + // Reserve room for the generation stamps bracketing the payload. + if bytes.len() + FRAME_OVERHEAD > self.slot_bytes { + return None; + } + let slot = self.free.lock().pop()?; + let addr = self.base + slot as usize * self.slot_bytes; + // A fresh, monotonically increasing stamp for THIS lease; a later lease of the + // same slot gets a strictly larger gen. Relaxed is fine: we only need + // uniqueness, and the cross-process happens-before is the descriptor ship. + let gen = self.gen.fetch_add(1, Ordering::Relaxed); + let gen_le = gen.to_le_bytes(); + // SAFETY: `slot` is exclusively leased (popped from `free`, returned only via + // free_room); [addr, addr + FRAME_OVERHEAD + len) is within the arena + // (len + FRAME_OVERHEAD <= slot_bytes) and disjoint from every other leased + // slot; the worker reads it only after we ship the descriptor below, and the + // TTL keeps it leased past the worker's max READ window. The arena is raw, + // leaked memory (no aliasing typed reference). So there is no data race and no + // overlap. Write the header gen FIRST and the trailer gen LAST so a racing + // reuse is visible at one stamp or the other (seqlock framing). + #[allow(unsafe_code)] + unsafe { + let p = addr as *mut u8; + std::ptr::copy_nonoverlapping(gen_le.as_ptr(), p, GEN_BYTES); + std::ptr::copy_nonoverlapping(bytes.as_ptr(), p.add(GEN_BYTES), bytes.len()); + std::ptr::copy_nonoverlapping( + gen_le.as_ptr(), + p.add(GEN_BYTES + bytes.len()), + GEN_BYTES, + ); + } + self.occupied.insert(room, OccSlot { slot, at: now }); + Some((slot, addr as u64, gen)) + } + + /// Return the slot leased for `room` to the free list (idempotent). + fn free_room(&self, room: i64) { + if let Some((_room, occ)) = self.occupied.remove(&room) { + self.free.lock().push(occ.slot); + } + } + + /// Reclaim every slot whose lease is older than `ttl` as of `now` (the lost-notif + /// net). Returns the count reclaimed. `now`/`ttl` are parameters so the reclaim + /// policy is testable without sleeping. + fn reap_stale(&self, now: Instant, ttl: Duration) -> usize { + let stale: Vec = self + .occupied + .iter() + .filter(|e| now.duration_since(e.value().at) >= ttl) + .map(|e| *e.key()) + .collect(); + let n = stale.len(); + for room in stale { + self.free_room(room); + } + n + } +} + +/// Pre-registered host-DRAM arena: a `SlotPool` plus the single persistent NIXL +/// registration that keeps the backing memory pinned for the worker's one-sided READ. +/// One registration (`_handle`) covers the whole arena; each image's pixels are +/// memcpy'd into a leased slot. A slot is exclusively owned between lease and free, +/// and the write happens-before the descriptor ship which happens-before the worker's +/// READ, so no slot is ever written and read concurrently. The hot path touches no +/// NIXL agent state (only the free-list lock + a lockless memcpy into owned memory). +struct SlotArena { + pool: SlotPool, + /// The single persistent registration; kept alive for the arena's life. + _handle: RegistrationHandle, +} + +/// Persistent gateway NIXL agent (one per process). The agent is touched only at +/// init (register the arena) and by the reaper (drain notifs); the hot path is +/// agent-free, so it stays behind a coarse Mutex without contending the fast path. +struct GatewayRdma { + agent: Agent, + // The UCX backend. We must pass it explicitly to register_memory (OptArgs) so the + // arena gets a UCX/rc rkey: register_memory(None) leaves the registration not bound + // to UCX, which is fine for TCP (copy) but makes a one-sided rc READ hang (no rkey). + backend: Backend, +} + +static AGENT: OnceLock>> = OnceLock::new(); +static ARENA: OnceLock> = OnceLock::new(); + +/// Lazily build the persistent agent (UCX backend) on first use; spawn the reaper. +/// Returns None (and callers fall back to inline) if NIXL init fails. +fn agent() -> Option<&'static Mutex> { + AGENT + .get_or_init(|| { + if !rdma_enabled() { + return None; + } + match init_agent() { + Ok(g) => { + spawn_reaper(); + debug!("EPD RDMA: gateway NIXL agent + UCX backend up"); + Some(Mutex::new(g)) + } + Err(e) => { + error!(error = ?e, "EPD RDMA: agent init failed; falling back to inline pixels"); + None + } + } + }) + .as_ref() +} + +/// Lazily build + register the slot arena (once) on first export. Requires the agent. +fn arena() -> Option<&'static SlotArena> { + ARENA + .get_or_init(|| { + let g = agent()?; + match build_arena(g) { + Ok(a) => { + debug!( + slots = a.pool.n_slots, + slot_bytes = a.pool.slot_bytes, + ttl_s = slot_ttl().as_secs(), + "EPD RDMA: pixel arena registered" + ); + Some(a) + } + Err(e) => { + error!(error = ?e, "EPD RDMA: arena registration failed; inline fallback"); + None + } + } + }) + .as_ref() +} + +fn build_arena(g: &Mutex) -> Result { + let n_slots = pool_slots(); + let slot_bytes = slot_bytes(); + let total = n_slots.saturating_mul(slot_bytes); + // Leak the arena for the process lifetime: `base` must stay registered + stable, + // and the gateway agent never shuts down. Holding it as raw memory (not a typed + // Box) is what makes the per-slot memcpy through `base` sound. + let boxed = vec![0u8; total].into_boxed_slice(); + let base = Box::into_raw(boxed) as *mut u8 as usize; + let region = ArenaRegion { base, size: total }; + let handle = { + let guard = g.lock(); + // Bind the registration to the UCX backend so it gets an rc rkey (needed for + // the worker's one-sided RDMA READ; without it the READ hangs over rc). + let mut opt = OptArgs::new()?; + opt.add_backend(&guard.backend)?; + guard.agent.register_memory(®ion, Some(&opt))? + }; + Ok(SlotArena { + pool: SlotPool::new(base, slot_bytes, n_slots), + _handle: handle, + }) +} + +fn init_agent() -> Result { + // enable_listen_thread + a fixed listen_port so the encode worker (initiator) + // can do the bidirectional NIXL metadata exchange (fetch_remote_metadata + + // send_local_metadata) against this gateway (target). Without the listener the + // worker's connect hits the ephemeral worker port -> "Connection refused" + // cross-node (same-host worked over shm). enable_prog_thread (default) keeps the + // UCX worker progressing so one-sided READs complete without our intervention. + let cfg = AgentConfig { + enable_listen_thread: true, + listen_port: i32::from(gw_listen_port()), + ..Default::default() + }; + let agent = Agent::new_configured(GATEWAY_AGENT_NAME, &cfg)?; + let (_mems, params) = agent.get_plugin_params("UCX")?; + let backend = agent.create_backend("UCX", ¶ms)?; + Ok(GatewayRdma { agent, backend }) +} + +/// Fixed agent name the encode worker passes to fetch_remote_metadata. +const GATEWAY_AGENT_NAME: &str = "smg-gateway-encode"; + +/// The gateway's RDMA listener IP (its RoCE address, e.g. 172.16.1.80) that the +/// encode worker dials for the metadata exchange. Empty -> RDMA disabled (inline). +fn gw_listen_ip() -> &'static str { + static IP: OnceLock = OnceLock::new(); + IP.get_or_init(|| std::env::var("SMG_RDMA_LISTEN_IP").unwrap_or_default()) +} + +/// The gateway's NIXL listener port (default 18515). +fn gw_listen_port() -> u16 { + static PORT: OnceLock = OnceLock::new(); + *PORT.get_or_init(|| { + std::env::var("SMG_RDMA_LISTEN_PORT") + .ok() + .and_then(|s| s.parse().ok()) + .unwrap_or(18515) + }) +} + +/// Arena slot count (`SMG_RDMA_POOL_SLOTS`, default 64). +fn pool_slots() -> usize { + static N: OnceLock = OnceLock::new(); + *N.get_or_init(|| { + std::env::var("SMG_RDMA_POOL_SLOTS") + .ok() + .and_then(|s| s.parse().ok()) + .filter(|&n| n > 0) + .unwrap_or(DEFAULT_POOL_SLOTS) + }) +} + +/// Per-slot byte capacity (`SMG_RDMA_SLOT_BYTES`, default 32 MiB). +fn slot_bytes() -> usize { + static B: OnceLock = OnceLock::new(); + *B.get_or_init(|| { + std::env::var("SMG_RDMA_SLOT_BYTES") + .ok() + .and_then(|s| s.parse().ok()) + .filter(|&n| n > 0) + .unwrap_or(DEFAULT_SLOT_BYTES) + }) +} + +/// Background reaper: drain free-notifs (tag = bootstrap_room) to return the leased +/// slot to the free list, plus a TTL sweep so a lost notif can never leak a slot. +fn spawn_reaper() { + std::thread::Builder::new() + .name("epd-rdma-reaper".into()) + .spawn(|| loop { + std::thread::sleep(REAPER_TICK); + let (Some(g), Some(a)) = ( + AGENT.get().and_then(|o| o.as_ref()), + ARENA.get().and_then(|o| o.as_ref()), + ) else { + continue; + }; + if let Ok(mut notifs) = NotificationMap::new() { + { + let guard = g.lock(); + let _ = guard.agent.get_notifications(&mut notifs, None); + } + if let Ok(map) = notifs.take_notifs() { + for (_agent, tags) in map { + for tag in tags { + if let Ok(room) = tag.parse::() { + a.pool.free_room(room); + } + } + } + } + } + // TTL sweep: reclaim slots whose READ-notif never arrived. `slot_ttl()` is + // derived to exceed the worker's max hold, so this never races a live READ. + let _ = a.pool.reap_stale(Instant::now(), slot_ttl()); + }) + .ok(); +} + +/// Stage `bytes` (the serialized pixel buffer for one image) into a pre-registered +/// arena slot keyed by `room`, and return the wire descriptor for the puller: +/// `[magic 8B][slot_addr u64 LE][gen u64 LE][room i64 LE][port u16 LE]` +/// `[listener_ip utf8]`. The worker fetches the gateway's (fixed) metadata once, +/// connects to the listener (ip:port), then READs `nbytes + FRAME_OVERHEAD` from +/// `slot_addr` and re-checks the stamps against `gen`. The slot is returned to +/// the free list on the worker's free-notif or the TTL. On any failure (no listener +/// IP, NIXL init, oversized image, or no free slot) returns `Err(bytes)` so the +/// caller re-attaches them as the inline payload (no behaviour change). +pub(crate) fn export_pixel_buffer(room: i64, bytes: Vec) -> Result, Vec> { + if gw_listen_ip().is_empty() { + // No listener IP configured -> cannot do the cross-node metadata exchange. + return Err(bytes); + } + let Some(arena) = arena() else { + return Err(bytes); + }; + let Some((slot, addr, gen)) = arena.pool.lease_and_write(room, &bytes, Instant::now()) else { + // Image too big for a slot, or the pool is momentarily exhausted. + return Err(bytes); + }; + + let ip = gw_listen_ip().as_bytes(); + let port = gw_listen_port(); + let mut descriptor = + Vec::with_capacity(DESCRIPTOR_MAGIC.len() + 8 + GEN_BYTES + 8 + 2 + ip.len()); + descriptor.extend_from_slice(DESCRIPTOR_MAGIC); + descriptor.extend_from_slice(&addr.to_le_bytes()); + descriptor.extend_from_slice(&gen.to_le_bytes()); + descriptor.extend_from_slice(&room.to_le_bytes()); + descriptor.extend_from_slice(&port.to_le_bytes()); + descriptor.extend_from_slice(ip); + debug!( + room, + addr, + slot, + "EPD RDMA: staged pixel slot (listener {}:{})", + gw_listen_ip(), + port + ); + Ok(descriptor) +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Build a `SlotPool` over a real (leaked) heap arena -- no NIXL agent, no RDMA. + /// Leaking matches production (the arena is process-lifetime) and keeps the raw + /// `base` valid for the whole test. + fn test_pool(n_slots: usize, slot_bytes: usize) -> SlotPool { + let buf = vec![0u8; n_slots * slot_bytes].into_boxed_slice(); + let base = Box::into_raw(buf) as *mut u8 as usize; + SlotPool::new(base, slot_bytes, n_slots) + } + + /// Read back what physically sits at a slot address (what a one-sided READ to that + /// address would return). + fn bytes_at(addr: u64, len: usize) -> Vec { + #[allow(unsafe_code)] + unsafe { + std::slice::from_raw_parts(addr as *const u8, len).to_vec() + } + } + + /// Reproduces the bug: a too-short TTL reclaims a slot while the worker is still + /// within its hold window, the next image re-leases the SAME address, and a late + /// READ against the old descriptor returns the WRONG image's bytes. + #[test] + fn short_ttl_recycles_slot_under_a_live_descriptor() { + let pool = test_pool(1, 64); + let t0 = Instant::now(); + let img_a = b"AAAAAAAAAAAAAAAA"; + let (_slot_a, addr_a, _gen_a) = pool.lease_and_write(1, img_a, t0).expect("lease A"); + + // Worker has NOT issued its READ yet (e.g. parked on landing-ring backpressure). + // The reaper runs with a 30s TTL after 31s -> reclaims the still-referenced slot. + let freed = pool.reap_stale(t0 + Duration::from_secs(31), Duration::from_secs(30)); + assert_eq!( + freed, 1, + "30s TTL reclaims a slot still under a live descriptor" + ); + + // The next image re-leases the freed slot -> same physical address. + let img_b = b"BBBBBBBBBBBBBBBB"; + let (_slot_b, addr_b, _gen_b) = pool + .lease_and_write(2, img_b, t0 + Duration::from_secs(31)) + .expect("lease B"); + assert_eq!(addr_a, addr_b, "freed slot's address is reused"); + + // The late READ for room 1 reads addr_a, whose payload (at +GEN_BYTES) now holds + // room 2's pixels. The TTL no longer protects this; the gen stamps do (next test). + assert_eq!( + bytes_at(addr_a + GEN_BYTES as u64, img_b.len()), + img_b, + "stale READ returns the WRONG image's payload (physical cross-wire remains)" + ); + } + + /// The gen guard: a recycled slot carries a fresh gen, so a worker validating the + /// slot's header/trailer stamps against its (older) descriptor gen detects the + /// cross-wire and fails the room -- independent of the TTL value. + #[test] + fn gen_guard_detects_recycled_slot() { + let pool = test_pool(1, 64); + let t0 = Instant::now(); + let img_a = b"AAAAAAAA"; + let (_a, addr_a, gen_a) = pool.lease_and_write(1, img_a, t0).expect("lease A"); + + // Recycle under a too-short TTL (the residual the gen guard backstops). + pool.reap_stale(t0 + Duration::from_secs(31), Duration::from_secs(30)); + let img_b = b"BBBBBBBB"; + let (_b, addr_b, gen_b) = pool + .lease_and_write(2, img_b, t0 + Duration::from_secs(31)) + .expect("lease B"); + assert_eq!(addr_a, addr_b, "slot address is physically reused"); + assert_ne!(gen_a, gen_b, "each lease gets a fresh, larger generation"); + + // The slot's header+trailer stamps now read gen_b. A worker holding room 1's + // descriptor (gen_a) compares against these and sees a mismatch -> reject. + let hdr = u64::from_le_bytes(bytes_at(addr_a, GEN_BYTES).try_into().unwrap()); + let trl = u64::from_le_bytes( + bytes_at(addr_a + (GEN_BYTES + img_b.len()) as u64, GEN_BYTES) + .try_into() + .unwrap(), + ); + assert_eq!(hdr, gen_b); + assert_eq!(trl, gen_b); + assert_ne!( + hdr, gen_a, + "stamp != descriptor gen -> worker fails the room" + ); + } + + /// A normal lease frames the payload as `[gen][payload][gen]` with both stamps equal + /// to the returned (and shipped-in-descriptor) gen. + #[test] + fn gen_framing_brackets_the_payload() { + let pool = test_pool(1, 64); + let t0 = Instant::now(); + let (_s, addr, gen) = pool.lease_and_write(7, b"PIXELS!!", t0).expect("lease"); + assert_eq!( + u64::from_le_bytes(bytes_at(addr, GEN_BYTES).try_into().unwrap()), + gen, + "header stamp == gen" + ); + assert_eq!(bytes_at(addr + GEN_BYTES as u64, 8), b"PIXELS!!", "payload"); + assert_eq!( + u64::from_le_bytes( + bytes_at(addr + (GEN_BYTES + 8) as u64, GEN_BYTES) + .try_into() + .unwrap() + ), + gen, + "trailer stamp == gen" + ); + } + + /// The fix invariant: the derived TTL strictly exceeds the worker's max hold, so + /// the reaper can never reclaim a slot the worker could still be reading. Fails + /// under the old flat `const SLOT_TTL = 30s`. + #[test] + fn slot_ttl_exceeds_worker_max_hold() { + assert!( + slot_ttl() > worker_max_hold(), + "slot_ttl {:?} must exceed worker_max_hold {:?} or a late READ cross-wires", + slot_ttl(), + worker_max_hold() + ); + } + + /// With the derived TTL, a reap at the worker's max hold does NOT reclaim, so the + /// address cannot be reused under a live descriptor: a second image finds the pool + /// full and the gateway falls back to inline (Err), never aliasing the slot. + #[test] + fn derived_ttl_keeps_slot_through_worker_hold() { + let pool = test_pool(1, 64); + let t0 = Instant::now(); + let (_slot_a, addr_a, _gen_a) = pool + .lease_and_write(1, b"AAAAAAAAAAAAAAAA", t0) + .expect("lease A"); + + let freed = pool.reap_stale(t0 + worker_max_hold(), slot_ttl()); + assert_eq!( + freed, 0, + "derived TTL must not reclaim within the worker hold" + ); + + // Slot still leased -> no free slot -> caller goes inline, no address reuse. + assert!( + pool.lease_and_write(2, b"BBBBBBBBBBBBBBBB", t0 + worker_max_hold()) + .is_none(), + "leased slot is not handed to a second image" + ); + assert_eq!( + bytes_at(addr_a + GEN_BYTES as u64, 16), + b"AAAAAAAAAAAAAAAA", + "the slot still holds room 1's pixels (payload at +GEN_BYTES) for its READ" + ); + } + + /// The recycle is not a one-slot artifact: with two slots, both reclaim under a + /// short TTL and a later lease reuses one of the freed addresses. + #[test] + fn multi_slot_addresses_recycle_after_reclaim() { + let pool = test_pool(2, 64); + let t0 = Instant::now(); + let (_a, addr_a, _ga) = pool.lease_and_write(1, b"A", t0).unwrap(); + let (_b, addr_b, _gb) = pool.lease_and_write(2, b"B", t0).unwrap(); + assert_ne!(addr_a, addr_b, "distinct leases get distinct addresses"); + + let freed = pool.reap_stale(t0 + Duration::from_secs(31), Duration::from_secs(30)); + assert_eq!(freed, 2); + + let (_c, addr_c, _gc) = pool + .lease_and_write(3, b"C", t0 + Duration::from_secs(31)) + .unwrap(); + assert!( + addr_c == addr_a || addr_c == addr_b, + "a reclaimed address is reused by the next lease" + ); + } + + /// An oversized image never leases (caller -> inline), and a normal free returns + /// the slot for reuse. + #[test] + fn oversize_rejected_and_free_returns_slot() { + // Slot holds FRAME_OVERHEAD (16) of gen stamps + payload, so 24 bytes => 8 of + // usable payload. + let pool = test_pool(1, FRAME_OVERHEAD + 8); + let t0 = Instant::now(); + assert!( + pool.lease_and_write(1, &[0u8; 9], t0).is_none(), + "payload + gen frame larger than a slot is rejected" + ); + let (_s, _addr, _g) = pool + .lease_and_write(1, &[1u8; 8], t0) + .expect("fits exactly"); + assert!( + pool.lease_and_write(2, &[2u8; 8], t0).is_none(), + "pool now full" + ); + pool.free_room(1); + assert!( + pool.lease_and_write(2, &[2u8; 8], t0).is_some(), + "freed slot is reusable" + ); + } +} diff --git a/model_gateway/src/routers/grpc/mm_rdma/stub.rs b/model_gateway/src/routers/grpc/mm_rdma/stub.rs new file mode 100644 index 0000000000..b3b54cc56f --- /dev/null +++ b/model_gateway/src/routers/grpc/mm_rdma/stub.rs @@ -0,0 +1,9 @@ +//! No-op multimodal RDMA transport used when the `mm-rdma` feature is disabled. + +pub(crate) fn rdma_enabled() -> bool { + false +} + +pub(crate) fn export_pixel_buffer(_room: i64, bytes: Vec) -> Result, Vec> { + Err(bytes) +} diff --git a/model_gateway/src/routers/grpc/mod.rs b/model_gateway/src/routers/grpc/mod.rs index 0948aa49e3..b3d1c7d3a5 100644 --- a/model_gateway/src/routers/grpc/mod.rs +++ b/model_gateway/src/routers/grpc/mod.rs @@ -7,6 +7,7 @@ pub(crate) mod common; pub(crate) mod context; pub(crate) mod epd_encode; pub(crate) mod harmony; +pub(crate) mod mm_rdma; pub(crate) mod multimodal; pub(crate) mod pd_router; // Used by routers/factory pub(crate) mod pipeline; diff --git a/model_gateway/src/routers/grpc/multimodal/config.rs b/model_gateway/src/routers/grpc/multimodal/config.rs index 55ca26fb82..d87c3e5634 100644 --- a/model_gateway/src/routers/grpc/multimodal/config.rs +++ b/model_gateway/src/routers/grpc/multimodal/config.rs @@ -11,6 +11,8 @@ use llm_multimodal::{ }; use tracing::{debug, warn}; +use super::pixel_cache::{pixel_cache_from_env, PixelCache}; + /// Cached model configuration files loaded from the tokenizer directory. #[derive(Debug, Clone)] pub(crate) struct MultimodalModelConfig { @@ -202,6 +204,8 @@ pub(crate) struct MultimodalComponents { pub model_registry: Arc, /// Shared reference to the app-level multimodal config cache. pub config_registry: Arc, + /// Optional host-DRAM cache of preprocessed per-image encoder inputs. + pub pixel_cache: Option>, } impl MultimodalComponents { @@ -220,6 +224,7 @@ impl MultimodalComponents { vision_processor_registry: Arc::new(VisionProcessorRegistry::with_defaults()), model_registry: Arc::new(ModelRegistry::default()), config_registry, + pixel_cache: pixel_cache_from_env(), }) } } diff --git a/model_gateway/src/routers/grpc/multimodal/mod.rs b/model_gateway/src/routers/grpc/multimodal/mod.rs index e3f69c5fb3..261755af4b 100644 --- a/model_gateway/src/routers/grpc/multimodal/mod.rs +++ b/model_gateway/src/routers/grpc/multimodal/mod.rs @@ -25,6 +25,7 @@ use llm_multimodal::{ mod assemble; mod config; mod detect; +mod pixel_cache; mod process; mod serialize; mod transport; diff --git a/model_gateway/src/routers/grpc/multimodal/pixel_cache.rs b/model_gateway/src/routers/grpc/multimodal/pixel_cache.rs new file mode 100644 index 0000000000..0ce3e7e358 --- /dev/null +++ b/model_gateway/src/routers/grpc/multimodal/pixel_cache.rs @@ -0,0 +1,256 @@ +//! Host-DRAM LRU cache of preprocessed per-image encoder inputs for the gateway +//! multimodal path. Disabled by default (`SMG_MM_PIXEL_CACHE_MB` unset / 0). + +use std::{ + mem::size_of, + sync::{Arc, OnceLock}, +}; + +use llm_multimodal::{ModelSpecificValue, PreprocessedEncoderInputs}; +use lru::LruCache; +use parking_lot::Mutex; + +/// Identifies a preprocessed image output. Hashing raw image bytes alone is not +/// enough because the same bytes preprocess differently under another model config. +#[derive(Clone, PartialEq, Eq, Hash, Debug)] +pub(crate) struct PixelCacheKey { + /// blake3 hex digest of the raw encoded image bytes (`ImageFrame.hash`). + pub image_hash: String, + /// Stable hash of model identity/config for this deployment. + pub config_fingerprint: u64, +} + +/// Cached per-image vision processor output before backend-specific serialization. +#[derive(Debug)] +pub(crate) struct CachedPreprocessedItem { + pub preprocessed: PreprocessedEncoderInputs, +} + +impl CachedPreprocessedItem { + fn heap_bytes(&self) -> usize { + preprocessed_heap_bytes(&self.preprocessed) + } +} + +fn preprocessed_heap_bytes(preprocessed: &PreprocessedEncoderInputs) -> usize { + let model_specific: usize = preprocessed + .model_specific + .iter() + .map(|(key, value)| key.len() + model_specific_value_heap_bytes(value)) + .sum(); + preprocessed.encoder_input.len() * size_of::() + + preprocessed.encoder_input.ndim() * size_of::() + + preprocessed.feature_token_counts.len() * size_of::() + + preprocessed.item_sizes.len() * size_of::<(u32, u32)>() + + model_specific +} + +fn model_specific_value_heap_bytes(value: &ModelSpecificValue) -> usize { + match value { + ModelSpecificValue::Tensor { data, shape } => { + data.len() * size_of::() + shape.len() * size_of::() + } + ModelSpecificValue::IntTensor { data, shape } => { + data.len() * size_of::() + shape.len() * size_of::() + } + ModelSpecificValue::UintTensor { data, shape } => { + data.len() * size_of::() + shape.len() * size_of::() + } + ModelSpecificValue::Int(_) => size_of::(), + ModelSpecificValue::Float(_) => size_of::(), + ModelSpecificValue::IntVec(values) => values.len() * size_of::(), + ModelSpecificValue::UintVec(values) => values.len() * size_of::(), + ModelSpecificValue::FloatVec(values) => values.len() * size_of::(), + ModelSpecificValue::TupleVec(values) => values.len() * size_of::<(u32, u32)>(), + ModelSpecificValue::Bool(_) => size_of::(), + } +} + +const PIXEL_CACHE_KEY_OVERHEAD: usize = 128; + +struct PixelCacheInner { + map: LruCache>, + cur_bytes: usize, + max_bytes: usize, +} + +/// Thread-safe, byte-budgeted LRU of per-image vision processor outputs. +pub(crate) struct PixelCache { + inner: Mutex, +} + +impl PixelCache { + pub(crate) fn new(max_bytes: usize) -> Self { + Self { + inner: Mutex::new(PixelCacheInner { + map: LruCache::unbounded(), + cur_bytes: 0, + max_bytes, + }), + } + } + + pub(crate) fn get(&self, key: &PixelCacheKey) -> Option> { + let mut inner = self.inner.lock(); + inner.map.get(key).cloned() + } + + pub(crate) fn insert(&self, key: PixelCacheKey, value: Arc) { + let entry_bytes = value.heap_bytes() + PIXEL_CACHE_KEY_OVERHEAD; + let mut inner = self.inner.lock(); + if entry_bytes > inner.max_bytes { + return; + } + let PixelCacheInner { + map, + cur_bytes, + max_bytes, + } = &mut *inner; + if let Some(previous) = map.put(key, value) { + *cur_bytes = cur_bytes.saturating_sub(previous.heap_bytes() + PIXEL_CACHE_KEY_OVERHEAD); + } + *cur_bytes += entry_bytes; + while *cur_bytes > *max_bytes { + match map.pop_lru() { + Some((_, evicted)) => { + *cur_bytes = + cur_bytes.saturating_sub(evicted.heap_bytes() + PIXEL_CACHE_KEY_OVERHEAD); + } + None => break, + } + } + } + + #[cfg(test)] + fn current_bytes(&self) -> usize { + self.inner.lock().cur_bytes + } + + #[cfg(test)] + fn len(&self) -> usize { + self.inner.lock().map.len() + } +} + +pub(crate) fn pixel_cache_from_env() -> Option> { + static CACHE: OnceLock>> = OnceLock::new(); + CACHE + .get_or_init(|| { + let mb = std::env::var("SMG_MM_PIXEL_CACHE_MB") + .ok() + .and_then(|raw| raw.trim().parse::().ok()) + .unwrap_or(0); + if mb == 0 { + return None; + } + let max_bytes = mb.saturating_mul(1024 * 1024); + tracing::info!( + target: "smg::request", + cache_mb = mb, + "multimodal pixel_values cache enabled (host DRAM, preprocessed encoder inputs)" + ); + Some(Arc::new(PixelCache::new(max_bytes))) + }) + .clone() +} + +pub(crate) fn config_fingerprint(tokenizer_id: &str, config: &serde_json::Value) -> u64 { + let mut hasher = blake3::Hasher::new(); + hasher.update(tokenizer_id.as_bytes()); + hasher.update(b"\0"); + hasher.update(config.to_string().as_bytes()); + let digest = hasher.finalize(); + let mut head = [0u8; 8]; + head.copy_from_slice(&digest.as_bytes()[..8]); + u64::from_le_bytes(head) +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use ndarray::{ArrayD, IxDyn}; + + use super::*; + + fn pixel_cache_item(token_count: usize, payload: usize) -> Arc { + Arc::new(CachedPreprocessedItem { + preprocessed: PreprocessedEncoderInputs { + encoder_input: ArrayD::from_elem(IxDyn(&[1, payload]), token_count as f32), + feature_token_counts: vec![token_count], + item_sizes: vec![(payload as u32, 1)], + model_specific: HashMap::new(), + }, + }) + } + + fn pixel_cache_key(hash: &str) -> PixelCacheKey { + PixelCacheKey { + image_hash: hash.to_string(), + config_fingerprint: 7, + } + } + + #[test] + fn pixel_cache_hit_returns_shared_arc() { + let cache = PixelCache::new(1024 * 1024); + let value = pixel_cache_item(10, 64); + cache.insert(pixel_cache_key("a"), value.clone()); + let got = cache.get(&pixel_cache_key("a")).expect("hit"); + assert_eq!(got.preprocessed.feature_token_counts, vec![10]); + assert!(Arc::ptr_eq(&got, &value)); + assert!(cache.get(&pixel_cache_key("missing")).is_none()); + } + + #[test] + fn pixel_cache_key_distinguishes_fingerprint() { + let cache = PixelCache::new(1024 * 1024); + let mut k_other_model = pixel_cache_key("a"); + k_other_model.config_fingerprint = 99; + cache.insert(pixel_cache_key("a"), pixel_cache_item(1, 8)); + cache.insert(k_other_model.clone(), pixel_cache_item(3, 8)); + assert_eq!( + cache + .get(&pixel_cache_key("a")) + .unwrap() + .preprocessed + .feature_token_counts, + vec![1] + ); + assert_eq!( + cache + .get(&k_other_model) + .unwrap() + .preprocessed + .feature_token_counts, + vec![3] + ); + assert_eq!(cache.len(), 2); + } + + #[test] + fn pixel_cache_evicts_lru_when_over_budget() { + let entry = pixel_cache_item(1, 4096).heap_bytes() + PIXEL_CACHE_KEY_OVERHEAD; + let budget = entry * 2 + entry / 2; + let cache = PixelCache::new(budget); + cache.insert(pixel_cache_key("a"), pixel_cache_item(1, 4096)); + cache.insert(pixel_cache_key("b"), pixel_cache_item(2, 4096)); + assert!(cache.get(&pixel_cache_key("a")).is_some()); + assert!(cache.get(&pixel_cache_key("b")).is_some()); + assert!(cache.get(&pixel_cache_key("a")).is_some()); + cache.insert(pixel_cache_key("c"), pixel_cache_item(3, 4096)); + assert!(cache.get(&pixel_cache_key("a")).is_some()); + assert!(cache.get(&pixel_cache_key("b")).is_none()); + assert!(cache.get(&pixel_cache_key("c")).is_some()); + assert!(cache.current_bytes() <= budget); + } + + #[test] + fn pixel_cache_oversized_entry_is_bypassed() { + let cache = PixelCache::new(32); + cache.insert(pixel_cache_key("huge"), pixel_cache_item(1, 1024)); + assert!(cache.get(&pixel_cache_key("huge")).is_none()); + assert_eq!(cache.len(), 0); + assert_eq!(cache.current_bytes(), 0); + } +} diff --git a/model_gateway/src/routers/grpc/multimodal/process.rs b/model_gateway/src/routers/grpc/multimodal/process.rs index 194d2399c1..d9ece9d6fd 100644 --- a/model_gateway/src/routers/grpc/multimodal/process.rs +++ b/model_gateway/src/routers/grpc/multimodal/process.rs @@ -9,7 +9,8 @@ use std::{sync::Arc, time::Instant}; use anyhow::Result; use llm_multimodal::{ AsyncMultiModalTracker, ImageFrame, Modality, ModelMetadata, PlaceholderRange, - PreprocessedEncoderInputs, PromptReplacement, TrackedMedia, TrackerOutput, VideoClip, + PreProcessorConfig, PreprocessedEncoderInputs, PromptReplacement, TrackedMedia, TrackerOutput, + VideoClip, VisionProcessorRegistry, }; use llm_tokenizer::TokenizerTrait; use openai_protocol::{chat::ChatMessage, messages::InputMessage}; @@ -18,8 +19,9 @@ use tracing::{debug, info, warn}; use super::{ config::MultimodalComponents, detect::{extract_content_parts, extract_content_parts_messages}, - log_mm_timing_enabled, MultimodalIntermediate, MultimodalOutput, - PrecomputedMultimodalIntermediate, + log_mm_timing_enabled, + pixel_cache::{config_fingerprint, CachedPreprocessedItem, PixelCache, PixelCacheKey}, + MultimodalIntermediate, MultimodalOutput, PrecomputedMultimodalIntermediate, }; /// Resolve the placeholder token string for a multimodal model. @@ -233,77 +235,93 @@ async fn process_multimodal_parts( }; let preprocess_started = Instant::now(); - let registry = components.vision_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 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}")); + let preprocessed: PreprocessedEncoderInputs = if let (Some(cache), Modality::Image, 1) = + (components.pixel_cache.clone(), modality, images.len()) + { + preprocess_image_cached( + cache, + &images[0], + components.vision_processor_registry.clone(), + model_id.to_string(), + model_type.map(String::from), + pp_config, + config_fingerprint(tokenizer_id, &model_config.config), + ) + .await? + } else { + let registry = components.vision_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 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" - ); + 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" - ); + Err(error) => { + warn!( + error = %error, + "RGB video frame refs are invalid; falling back to materialized frames" + ); + } } } - } - 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 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}")) + } + _ => Err(anyhow::anyhow!( + "Unsupported modality for preprocessing: {modality}" + )), } - _ => Err(anyhow::anyhow!( - "Unsupported modality for preprocessing: {modality}" - )), - } - }) - .await - .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))??; + }) + .await + .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))??; + preprocessed + }; let preprocess_elapsed_ms = preprocess_started.elapsed().as_secs_f64() * 1000.0; debug!( @@ -405,6 +423,61 @@ async fn process_multimodal_parts( }) } +/// Pixel-cache image preprocessing for single-image requests. +async fn preprocess_image_cached( + cache: Arc, + image: &Arc, + registry: Arc, + model_id: String, + model_type: Option, + pp_config: PreProcessorConfig, + fingerprint: u64, +) -> Result { + let key = PixelCacheKey { + image_hash: image.hash.clone(), + config_fingerprint: fingerprint, + }; + if let Some(cached) = cache.get(&key) { + return Ok(cached.preprocessed.clone()); + } + + let preprocessed = preprocess_image_batch( + registry, + model_id, + model_type, + pp_config, + std::slice::from_ref(image), + ) + .await?; + cache.insert( + key, + Arc::new(CachedPreprocessedItem { + preprocessed: preprocessed.clone(), + }), + ); + Ok(preprocessed) +} + +async fn preprocess_image_batch( + registry: Arc, + model_id: String, + model_type: Option, + pp_config: PreProcessorConfig, + images: &[Arc], +) -> Result { + let raw_images: Vec = images.iter().map(|f| f.image.clone()).collect(); + tokio::task::spawn_blocking(move || { + let processor = registry + .find(&model_id, model_type.as_deref()) + .ok_or_else(|| anyhow::anyhow!("No vision processor found for model: {model_id}"))?; + processor + .preprocess(&raw_images, &pp_config) + .map_err(|e| anyhow::anyhow!("Image preprocessing failed: {e}")) + }) + .await + .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))? +} + /// Output of token expansion, containing both full structural and patch-only ranges. struct ExpandedTokens { /// The expanded token ID sequence. diff --git a/model_gateway/src/routers/grpc/multimodal/serialize.rs b/model_gateway/src/routers/grpc/multimodal/serialize.rs index 2c6c744e7b..bdf5028dab 100644 --- a/model_gateway/src/routers/grpc/multimodal/serialize.rs +++ b/model_gateway/src/routers/grpc/multimodal/serialize.rs @@ -489,6 +489,9 @@ mod tests { TokenSpeedTensorStorage::Inline(_) => { panic!("expected SHM at/above the threshold when /dev/shm is writable") } + TokenSpeedTensorStorage::Remote(_) => { + panic!("unexpected remote payload in the SHM threshold test") + } } } } diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index d6e169f1d3..f808e47801 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -19,6 +19,7 @@ use std::{ use futures_util::StreamExt; use memmap2::MmapOptions; +use rand::RngExt; #[cfg(target_os = "linux")] use rustix::fs::FallocateFlags; use smg_grpc_client::{ @@ -146,9 +147,9 @@ pub struct TensorBytes { pub dtype: String, } -/// TokenSpeed tensor with storage that can be either inline bytes or an -/// already-published SHM handle. The latter lets SMG serialize large encoder -/// inputs directly into SHM instead of building a temporary Vec first. +/// TokenSpeed tensor with an explicit payload transport. Inline, SHM, and +/// remote descriptors all converge here before becoming generated proto fields, +/// so stages do not need to mutate `TensorData` oneofs directly. #[derive(Debug, Clone)] pub struct TokenSpeedTensor { pub storage: TokenSpeedTensorStorage, @@ -160,6 +161,7 @@ pub struct TokenSpeedTensor { pub enum TokenSpeedTensorStorage { Inline(Vec), Shm(common::ShmHandle), + Remote(common::RemoteTensorHandle), } impl TokenSpeedTensor { @@ -179,10 +181,58 @@ impl TokenSpeedTensor { } } + pub fn remote(handle: common::RemoteTensorHandle, shape: Vec, dtype: String) -> Self { + Self { + storage: TokenSpeedTensorStorage::Remote(handle), + shape, + dtype, + } + } + + pub fn try_export_nixl_remote(self, room: i64) -> Self { + if !crate::routers::grpc::mm_rdma::rdma_enabled() { + return self; + } + + let Self { + storage, + shape, + dtype, + } = self; + let data = match storage { + TokenSpeedTensorStorage::Inline(data) => data, + storage => { + return Self { + storage, + shape, + dtype, + }; + } + }; + if data.is_empty() { + return Self::inline(data, shape, dtype); + } + + let nbytes = data.len() as u64; + match crate::routers::grpc::mm_rdma::export_pixel_buffer(room, data) { + Ok(descriptor) => Self::remote( + common::RemoteTensorHandle { + transport: "nixl".to_string(), + descriptor, + nbytes, + }, + shape, + dtype, + ), + Err(data) => Self::inline(data, shape, dtype), + } + } + pub fn nbytes(&self) -> usize { match &self.storage { TokenSpeedTensorStorage::Inline(data) => data.len(), TokenSpeedTensorStorage::Shm(handle) => handle.nbytes as usize, + TokenSpeedTensorStorage::Remote(handle) => handle.nbytes as usize, } } } @@ -288,6 +338,23 @@ impl TrtllmMultimodalData { } impl TokenSpeedMultimodalData { + /// Export inline encoder-input payloads over NIXL for normal TokenSpeed + /// Generate requests (single-worker and PD prefill legs). EPD encode uses + /// its own room-matched export path because the room must also be injected + /// into the encode->prefill handshake. + pub fn try_export_encoder_inputs_nixl_remote(mut self) -> Self { + if !crate::routers::grpc::mm_rdma::rdma_enabled() { + return self; + } + for item in &mut self.items { + let room = rand::rng().random_range(0..i64::MAX); + let placeholder = TokenSpeedTensor::inline(Vec::new(), Vec::new(), String::new()); + let encoder_input = std::mem::replace(&mut item.encoder_input, placeholder); + item.encoder_input = encoder_input.try_export_nixl_remote(room); + } + self + } + /// Convert to TokenSpeed proto MultimodalInputs. The EPD prefill leg drops /// each item's encoder_input afterward via `clear_mm_pixel_values`. pub fn into_proto(self) -> tokenspeed::MultimodalInputs { @@ -358,6 +425,10 @@ fn tokenspeed_tensor_to_proto( Metrics::record_mm_tensor("tokenspeed", "shm", handle.nbytes as usize); tokenspeed::tensor_data::Payload::Shm(handle) } + TokenSpeedTensorStorage::Remote(handle) => { + Metrics::record_mm_tensor("tokenspeed", "remote", handle.nbytes as usize); + tokenspeed::tensor_data::Payload::Remote(handle) + } }; tokenspeed::TensorData { @@ -2098,6 +2169,43 @@ mod tests { assert!(!tokenspeed_shm_path(&handle.name).exists()); } + #[test] + fn tokenspeed_remote_encoder_input_into_proto_uses_remote_payload() { + let proto = TokenSpeedMultimodalData { + items: vec![TokenSpeedMultimodalItem { + modality: TokenSpeedModality::Image, + encoder_input: TokenSpeedTensor::remote( + common::RemoteTensorHandle { + transport: "nixl".to_string(), + descriptor: vec![1, 2, 3], + nbytes: 8, + }, + vec![1, 2], + "bfloat16".to_string(), + ), + model_specific_tensors: HashMap::new(), + placeholder_token_id: Some(151655), + mm_placeholders: vec![(4, 2)], + content_hash: vec![7; 32], + }], + shm_enabled: true, + shm_min_bytes: 0, + } + .into_proto(); + + let tensor = proto.items[0].encoder_input.as_ref().unwrap(); + assert_eq!(tensor.shape, vec![1, 2]); + assert_eq!(tensor.dtype, "bfloat16"); + match tensor.payload.as_ref() { + Some(tokenspeed::tensor_data::Payload::Remote(handle)) => { + assert_eq!(handle.transport, "nixl"); + assert_eq!(handle.descriptor, vec![1, 2, 3]); + assert_eq!(handle.nbytes, 8); + } + _ => panic!("expected remote TensorData payload"), + } + } + fn inline_tensor_data(tensor: &tokenspeed::TensorData) -> &[u8] { match tensor.payload.as_ref() { Some(tokenspeed::tensor_data::Payload::Inline(data)) => data, From 18079f30fcfecfaf700c4b54dfd432f684c3b712 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Wed, 8 Jul 2026 18:18:04 -0700 Subject: [PATCH 2/4] refactor(multimodal): make RDMA pixel transport a first-class TransportMode Fold the multimodal pixel RDMA lane into the engine-agnostic transport config layer instead of a standalone env gate. Add `TransportMode::Rdma` so it is selectable via `--multimodal-tensor-transport rdma` / `SMG_MM_TENSOR_TRANSPORT=rdma` alongside inline/shm/auto; the RDMA gate now consults the resolved transport mode, keeping the legacy `SMG_MM_PIXEL_RDMA` env as a backward-compatible fallback. - openai-protocol: TransportMode::Rdma variant + parse/as_str/Display/FromStr - transport: resolve_mm_shm_enabled treats rdma as non-SHM; feature-gated mm_default_transport_is_rdma() accessor the gate consults - mm_rdma: rdma_enabled() honors TransportMode::Rdma (mm-rdma feature) - CLI value_parser, Python arg, binding, and docs recognize the rdma value Co-authored-by: chenht2022 Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- bindings/python/src/lib.rs | 2 +- bindings/python/src/smg/router_args.py | 4 ++-- crates/protocols/src/worker.rs | 10 ++++++++-- docs/reference/configuration.md | 2 +- model_gateway/src/config/types.rs | 2 +- model_gateway/src/main.rs | 2 +- model_gateway/src/routers/grpc/mm_rdma/nixl.rs | 16 +++++++++++----- .../src/routers/grpc/multimodal/mod.rs | 2 ++ .../src/routers/grpc/multimodal/transport.rs | 17 +++++++++++++++-- 9 files changed, 42 insertions(+), 15 deletions(-) diff --git a/bindings/python/src/lib.rs b/bindings/python/src/lib.rs index 7e4e58a64b..938a58063d 100644 --- a/bindings/python/src/lib.rs +++ b/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(), } }) }) diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index 213266ce81..cfb3218649 100644 --- a/bindings/python/src/smg/router_args.py +++ b/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 mode: inline (default), shm, auto, or rdma (NIXL pixel lane; requires the mm-rdma build)", ) parser.add_argument( f"--{prefix}multimodal-shm-min-bytes", diff --git a/crates/protocols/src/worker.rs b/crates/protocols/src/worker.rs index d13d421268..b1f605c67c 100644 --- a/crates/protocols/src/worker.rs +++ b/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 { 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::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")) } } diff --git a/docs/reference/configuration.md b/docs/reference/configuration.md index bd5648c0cb..ea9094e197 100644 --- a/docs/reference/configuration.md +++ b/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 diff --git a/model_gateway/src/config/types.rs b/model_gateway/src/config/types.rs index 44afa8f539..a17cc4bd1f 100644 --- a/model_gateway/src/config/types.rs +++ b/model_gateway/src/config/types.rs @@ -44,7 +44,7 @@ pub struct RouterConfig { /// observability from routing. #[serde(default)] pub engine_metrics: bool, - /// Global multimodal tensor transport mode (`inline` | `shm` | `auto`). + /// Global multimodal tensor transport mode (`inline` | `shm` | `auto` | `rdma`). /// Per-worker `WorkerSpec.multimodal_tensor_transport` overrides this; when /// unset, falls back to `SMG_MM_TENSOR_TRANSPORT`, then `inline`. #[serde(default, skip_serializing_if = "Option::is_none")] diff --git a/model_gateway/src/main.rs b/model_gateway/src/main.rs index 2b7da4887c..e2e8caa563 100644 --- a/model_gateway/src/main.rs +++ b/model_gateway/src/main.rs @@ -150,7 +150,7 @@ enum Commands { /// Parse the `--multimodal-tensor-transport` value into a `TransportMode`. fn parse_transport_mode(value: &str) -> Result { TransportMode::parse(value) - .ok_or_else(|| format!("invalid value '{value}'; expected inline, shm, or auto")) + .ok_or_else(|| format!("invalid value '{value}'; expected inline, shm, auto, or rdma")) } #[derive(Parser, Debug)] diff --git a/model_gateway/src/routers/grpc/mm_rdma/nixl.rs b/model_gateway/src/routers/grpc/mm_rdma/nixl.rs index f708b6a433..6da22f32c1 100644 --- a/model_gateway/src/routers/grpc/mm_rdma/nixl.rs +++ b/model_gateway/src/routers/grpc/mm_rdma/nixl.rs @@ -35,14 +35,20 @@ use nixl_sys::{ use parking_lot::Mutex; use tracing::{debug, error}; -/// `SMG_MM_PIXEL_RDMA` gate (cached in a process-wide `OnceLock`). +use crate::routers::grpc::multimodal::mm_default_transport_is_rdma; + +/// Whether the RDMA pixel lane is active. Selected by the first-class +/// `TransportMode::Rdma` (`--multimodal-tensor-transport rdma` / +/// `SMG_MM_TENSOR_TRANSPORT=rdma`), with the legacy `SMG_MM_PIXEL_RDMA` env as a +/// backward-compatible fallback. Cached in a process-wide `OnceLock`. pub(crate) fn rdma_enabled() -> bool { static ENABLED: OnceLock = OnceLock::new(); *ENABLED.get_or_init(|| { - matches!( - std::env::var("SMG_MM_PIXEL_RDMA").as_deref(), - Ok("1") | Ok("true") - ) + mm_default_transport_is_rdma() + || matches!( + std::env::var("SMG_MM_PIXEL_RDMA").as_deref(), + Ok("1") | Ok("true") + ) }) } diff --git a/model_gateway/src/routers/grpc/multimodal/mod.rs b/model_gateway/src/routers/grpc/multimodal/mod.rs index 261755af4b..b35fd0e9cf 100644 --- a/model_gateway/src/routers/grpc/multimodal/mod.rs +++ b/model_gateway/src/routers/grpc/multimodal/mod.rs @@ -43,6 +43,8 @@ pub(crate) use process::{ process_multimodal, process_multimodal_messages, resolve_placeholder_token, }; pub(crate) use transport::init_mm_transport_defaults; +#[cfg(feature = "mm-rdma")] +pub(crate) use transport::mm_default_transport_is_rdma; /// Whether verbose multimodal timing logs are enabled via `SMG_LOG_MM_TIMING`. /// Read from the environment once and cached; the flag is not expected to change diff --git a/model_gateway/src/routers/grpc/multimodal/transport.rs b/model_gateway/src/routers/grpc/multimodal/transport.rs index fdf923aa9c..4a801c2675 100644 --- a/model_gateway/src/routers/grpc/multimodal/transport.rs +++ b/model_gateway/src/routers/grpc/multimodal/transport.rs @@ -139,7 +139,8 @@ pub(super) fn resolve_mm_shm_enabled( TransportMode::Auto => { worker_shares_dev_shm(workers, skip_pixel_values) && mm_shm_dev_writable() } - TransportMode::Inline => false, + // `rdma` routes large tensors through the NIXL pixel lane, not SHM. + TransportMode::Inline | TransportMode::Rdma => false, } } @@ -149,6 +150,18 @@ pub(super) fn resolve_mm_shm_min_bytes(workers: Option<&WorkerSelection>) -> usi worker_shm_min_bytes_override(workers).unwrap_or_else(|| mm_transport_defaults().shm_min_bytes) } +/// Whether the router-level transport mode selects the RDMA pixel lane. +/// +/// The `mm_rdma` gate consults this so `rdma` is a first-class [`TransportMode`] +/// alongside `inline`/`shm`/`auto` (settable via `--multimodal-tensor-transport` +/// / `SMG_MM_TENSOR_TRANSPORT`). The legacy `SMG_MM_PIXEL_RDMA` env stays a +/// fallback inside the gate for backward compatibility. Only compiled with the +/// `mm-rdma` feature, since the no-op RDMA shim never consults it. +#[cfg(feature = "mm-rdma")] +pub(crate) fn mm_default_transport_is_rdma() -> bool { + mm_transport_defaults().mode == TransportMode::Rdma +} + fn worker_transport_mode_override(workers: Option<&WorkerSelection>) -> Option { primary_worker(workers)? .metadata() @@ -251,7 +264,7 @@ fn log_unknown_transport_once(value: &str) { WARNED.get_or_init(|| { warn!( value, - "Unknown multimodal tensor transport value; expected inline|shm|auto, using inline" + "Unknown multimodal tensor transport value; expected inline|shm|auto|rdma, using inline" ); }); } From 492198782562cc434ca5ed949ae22bedaa73a96f Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Wed, 8 Jul 2026 20:18:00 -0700 Subject: [PATCH 3/4] fix(multimodal): satisfy clippy and ruff after rebase Make #1853 pass CI (clippy --all-targets --all-features -D warnings, ruff): - absolute_paths: import crate::vision::scratch (kimi_k25, transforms) and crate::routers::grpc::mm_rdma (proto_wrapper) instead of absolute-path calls - allow_attributes: #[expect(unsafe_code, reason=...)] not #[allow] for the NIXL FFI unsafe blocks (nixl.rs) - ruff: shorten the rdma transport-mode help string under 120 cols Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- bindings/python/src/smg/router_args.py | 2 +- .../multimodal/src/vision/processors/kimi_k25.rs | 7 ++++--- crates/multimodal/src/vision/transforms.rs | 7 +++++-- model_gateway/src/routers/grpc/mm_rdma/nixl.rs | 15 ++++++++++++--- model_gateway/src/routers/grpc/proto_wrapper.rs | 8 +++++--- 5 files changed, 27 insertions(+), 12 deletions(-) diff --git a/bindings/python/src/smg/router_args.py b/bindings/python/src/smg/router_args.py index cfb3218649..4bbfb2c84d 100644 --- a/bindings/python/src/smg/router_args.py +++ b/bindings/python/src/smg/router_args.py @@ -573,7 +573,7 @@ def add_cli_args( type=str, choices=["inline", "shm", "auto", "rdma"], default=RouterArgs.multimodal_tensor_transport, - help="Multimodal tensor transport mode: inline (default), shm, auto, or rdma (NIXL pixel lane; requires the mm-rdma build)", + 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", diff --git a/crates/multimodal/src/vision/processors/kimi_k25.rs b/crates/multimodal/src/vision/processors/kimi_k25.rs index a2ea599a9e..1d9adf2bdf 100644 --- a/crates/multimodal/src/vision/processors/kimi_k25.rs +++ b/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}, }; @@ -158,7 +159,7 @@ impl KimiK25Processor { // 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 = crate::vision::scratch::take_f32(3 * canvas_pixels); + 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); @@ -264,7 +265,7 @@ impl VisionPreProcessor for KimiK25Processor { 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 = crate::vision::scratch::take_f32_cap(estimated_total); + let mut all_patches: Vec = scratch::take_f32_cap(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()); @@ -293,7 +294,7 @@ impl VisionPreProcessor for KimiK25Processor { // 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(); - crate::vision::scratch::give_f32(storage); + scratch::give_f32(storage); patches_per_image.push(num_patches as i64); } diff --git a/crates/multimodal/src/vision/transforms.rs b/crates/multimodal/src/vision/transforms.rs index 875773136c..ddb490f3c5 100644 --- a/crates/multimodal/src/vision/transforms.rs +++ b/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)] @@ -152,7 +155,7 @@ fn build_planar_tensor( ) -> Array3 { let pixels = h * w; // Pooled: this large per-image buffer is the data plane's hottest allocation. - let mut data = crate::vision::scratch::take_f32(3 * pixels); + 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); diff --git a/model_gateway/src/routers/grpc/mm_rdma/nixl.rs b/model_gateway/src/routers/grpc/mm_rdma/nixl.rs index 6da22f32c1..e797d623c8 100644 --- a/model_gateway/src/routers/grpc/mm_rdma/nixl.rs +++ b/model_gateway/src/routers/grpc/mm_rdma/nixl.rs @@ -135,7 +135,10 @@ struct ArenaRegion { impl MemoryRegion for ArenaRegion { // The NIXL trait declares this `unsafe`; the body just reinterprets a usize the // crate denies unsafe_code globally, so allow it for this trait impl. - #[allow(unsafe_code)] + #[expect( + unsafe_code, + reason = "NIXL MemoryRegion trait method is unsafe; body only reinterprets the stored base address as a pointer" + )] unsafe fn as_ptr(&self) -> *const u8 { self.base as *const u8 } @@ -217,7 +220,10 @@ impl SlotPool { // leaked memory (no aliasing typed reference). So there is no data race and no // overlap. Write the header gen FIRST and the trailer gen LAST so a racing // reuse is visible at one stamp or the other (seqlock framing). - #[allow(unsafe_code)] + #[expect( + unsafe_code, + reason = "seqlock gen-stamp write into the pre-registered NIXL slot arena" + )] unsafe { let p = addr as *mut u8; std::ptr::copy_nonoverlapping(gen_le.as_ptr(), p, GEN_BYTES); @@ -512,7 +518,10 @@ mod tests { /// Read back what physically sits at a slot address (what a one-sided READ to that /// address would return). fn bytes_at(addr: u64, len: usize) -> Vec { - #[allow(unsafe_code)] + #[expect( + unsafe_code, + reason = "read bytes back from a NIXL slot address for gen verification" + )] unsafe { std::slice::from_raw_parts(addr as *const u8, len).to_vec() } diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index f808e47801..9eb7362f2b 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -38,6 +38,8 @@ use smg_grpc_client::{ vllm_proto::{self as vllm, generate_complete::MatchedStop as VllmMatchedStop}, }; +use crate::routers::grpc::mm_rdma; + /// Backend-neutral encode->prefill bootstrap info for one multimodal item. /// /// Backend wrappers translate this into their own proto shape when supported. @@ -190,7 +192,7 @@ impl TokenSpeedTensor { } pub fn try_export_nixl_remote(self, room: i64) -> Self { - if !crate::routers::grpc::mm_rdma::rdma_enabled() { + if !mm_rdma::rdma_enabled() { return self; } @@ -214,7 +216,7 @@ impl TokenSpeedTensor { } let nbytes = data.len() as u64; - match crate::routers::grpc::mm_rdma::export_pixel_buffer(room, data) { + match mm_rdma::export_pixel_buffer(room, data) { Ok(descriptor) => Self::remote( common::RemoteTensorHandle { transport: "nixl".to_string(), @@ -343,7 +345,7 @@ impl TokenSpeedMultimodalData { /// its own room-matched export path because the room must also be injected /// into the encode->prefill handshake. pub fn try_export_encoder_inputs_nixl_remote(mut self) -> Self { - if !crate::routers::grpc::mm_rdma::rdma_enabled() { + if !mm_rdma::rdma_enabled() { return self; } for item in &mut self.items { From 04aeff1202e7713ea30aa10bb657f4eab2eb6a9a Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Wed, 8 Jul 2026 21:12:47 -0700 Subject: [PATCH 4/4] ci: re-trigger CI after infra fix Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>