From e5bc23003194f88e46dd33211b1b644eb2dd4a86 Mon Sep 17 00:00:00 2001 From: Simo Lin <25425177+slin1237@users.noreply.github.com> Date: Mon, 13 Jul 2026 09:10:53 -0700 Subject: [PATCH] refactor(multimodal): route RDMA through the shared tensor payload resolver Make RDMA a first-class transport in the engine-neutral tensor payload path (alongside inline/SHM), decouple it from TokenSpeed, and hoist the Python puller out of the tokenspeed package. vLLM emit stays gated off until its puller lands. Rust: - `resolve_mm_tensor_payload` takes an optional `MmRdmaExport` and resolves RDMA -> SHM -> inline; on export failure the bytes fall through so the payload is never dropped. Add `MmTensorPayload::Remote` + `Remote` arms to the vLLM/TokenSpeed payload mappers. - The general TokenSpeed export moves into `into_proto` (random slot key, encoder_input only); EPD stages explicitly with `bootstrap_room` via `stage_tokenspeed_tensor_rdma`. Delete `try_export_nixl_remote` and `try_export_encoder_inputs_nixl_remote`. - `VllmMultimodalData` carries `rdma_enabled` (held `false`: vLLM cannot pull RDMA payloads yet). Python: - Move `tokenspeed/rdma_pixel.py` to `smg_grpc_servicer/mm_rdma.py`, engine-neutral: `gateway_agent_name` is a constructor param and `feature_from_remote` returns the tensor, with the SHM publish moved to the TokenSpeed caller (dropping the `tokenspeed.runtime` imports). Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com> --- .../{tokenspeed/rdma_pixel.py => mm_rdma.py} | 41 ++-- .../tokenspeed/encoder_servicer.py | 29 ++- .../smg_grpc_servicer/tokenspeed/servicer.py | 5 +- model_gateway/src/routers/grpc/client.rs | 8 +- model_gateway/src/routers/grpc/epd_encode.rs | 15 +- .../src/routers/grpc/multimodal/assemble.rs | 2 + .../src/routers/grpc/proto_wrapper.rs | 212 +++++++++++------- 7 files changed, 184 insertions(+), 128 deletions(-) rename grpc_servicer/smg_grpc_servicer/{tokenspeed/rdma_pixel.py => mm_rdma.py} (92%) diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py b/grpc_servicer/smg_grpc_servicer/mm_rdma.py similarity index 92% rename from grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py rename to grpc_servicer/smg_grpc_servicer/mm_rdma.py index 2782cc9b80..44b7b17c47 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py +++ b/grpc_servicer/smg_grpc_servicer/mm_rdma.py @@ -1,4 +1,4 @@ -"""Shared NIXL pixel-payload puller for TokenSpeed multimodal inputs.""" +"""Engine-neutral NIXL pixel-payload puller for gateway-exported multimodal inputs.""" from __future__ import annotations @@ -15,7 +15,8 @@ _DESCRIPTOR_MAGIC = b"SMGRDMA1" _GEN_BYTES = 8 _FRAME_BYTES = 2 * _GEN_BYTES -_GATEWAY_AGENT_NAME = "smg-gateway-encode" +# The gateway's fixed NIXL agent name (see RDMA_GATEWAY_AGENT_NAME gateway-side). +DEFAULT_GATEWAY_AGENT_NAME = "smg-gateway-encode" _LOCAL_IP_CACHE: dict[str, bool] = {} @@ -102,8 +103,15 @@ def _parse_descriptor(td, explicit_room: int | None) -> _RemotePixelDescriptor: class RdmaPixelPuller: """Persistent NIXL READ agent for gateway-exported multimodal pixel buffers.""" - def __init__(self, *, agent_name: str, log_prefix: str): + def __init__( + self, + *, + agent_name: str, + log_prefix: str, + gateway_agent_name: str = DEFAULT_GATEWAY_AGENT_NAME, + ): self._log_prefix = log_prefix + self._gateway_agent_name = gateway_agent_name self._nixl_agent = None self._rdma_md_ready = set() self._rdma_md_lock = threading.Lock() @@ -171,7 +179,7 @@ def _ensure_remote_ready(self, ip: str, port: int, remote, room: int) -> None: return agent = self._nixl_agent - agent.fetch_remote_metadata(_GATEWAY_AGENT_NAME, ip, port) + agent.fetch_remote_metadata(self._gateway_agent_name, ip, port) send_md_env = os.environ.get("SMG_RDMA_SEND_MD") if send_md_env in ("1", "true"): @@ -194,7 +202,7 @@ def _ensure_remote_ready(self, ip: str, port: int, remote, room: int) -> None: ready = False for _ in range(5000): - if agent.check_remote_metadata(_GATEWAY_AGENT_NAME, remote): + if agent.check_remote_metadata(self._gateway_agent_name, remote): ready = True break time.sleep(0.001) @@ -202,15 +210,8 @@ def _ensure_remote_ready(self, ip: str, port: int, remote, room: int) -> None: 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)``.""" + def feature_from_remote(self, td, *, explicit_room: int | None, cast_to): + """Read a remote TensorData payload and return the feature tensor.""" import numpy as np import torch @@ -271,7 +272,7 @@ def feature_from_remote( "READ", local, remote, - _GATEWAY_AGENT_NAME, + self._gateway_agent_name, str(descriptor.room).encode(), ) read_deadline = time.monotonic() + float(os.environ.get("SMG_RDMA_READ_TIMEOUT_S", 60)) @@ -327,13 +328,7 @@ def feature_from_remote( 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 + # Copy out of the landing slot before it returns to the free ring. + return tensor if copied else tensor.clone() finally: self._landing_free.put(slot) diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py index 39b8f35cf0..a788bc0da1 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py @@ -22,7 +22,7 @@ tokenspeed_encoder_pb2_grpc, ) -from smg_grpc_servicer.tokenspeed.rdma_pixel import RdmaPixelPuller +from smg_grpc_servicer.mm_rdma import RdmaPixelPuller from smg_grpc_servicer.tokenspeed.servicer import TokenSpeedSchedulerServicer if TYPE_CHECKING: @@ -125,27 +125,24 @@ def _items_from_proto(self, mm_inputs, bootstrap_room: int = 0): # before the swap and pre-set on the item. td = item_proto.encoder_input if td.WhichOneof("payload") == "remote": - # EPD RDMA: pull the encoder input from 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 encoder tensor. - feature, feat_hash = self._rdma_pixel_puller.feature_from_remote( + feature = 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) + + # EPD_PIXEL_SHM publishes the feature to scheduler SHM so the ZMQ hop + # pickles ~KB instead of the full encoder tensor; the content hash is + # taken on the real bytes first. + 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) item = MultimodalDataItem( modality=item_modality, diff --git a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py index c631acdf64..6f1763f79c 100644 --- a/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py +++ b/grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py @@ -39,9 +39,9 @@ from tokenspeed.runtime.pd.kv_events import KVEventBatch from smg_grpc_servicer.kv_events import endpoint_for_rank, stream_kv_events +from smg_grpc_servicer.mm_rdma import RdmaPixelPuller 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 @@ -1072,11 +1072,10 @@ def _mm_inputs_from_itemized_proto( feature_started = time.perf_counter() if LOG_MM_TIMING else None encoder_input = item_proto.encoder_input if encoder_input.WhichOneof("payload") == "remote": - feature, _ = self._rdma_pixel_puller.feature_from_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) diff --git a/model_gateway/src/routers/grpc/client.rs b/model_gateway/src/routers/grpc/client.rs index 990a291d11..c541c4203f 100644 --- a/model_gateway/src/routers/grpc/client.rs +++ b/model_gateway/src/routers/grpc/client.rs @@ -534,9 +534,7 @@ impl GrpcClient { } Self::TokenSpeed(client) => { let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { - MultimodalData::TokenSpeed(data) => { - data.try_export_encoder_inputs_nixl_remote().into_proto() - } + MultimodalData::TokenSpeed(data) => data.into_proto(true), _ => unreachable!("caller guarantees matching variant"), }); finish_tokenspeed_request(tokenspeed_mm, |mm| { @@ -628,9 +626,7 @@ impl GrpcClient { } Self::TokenSpeed(client) => { let tokenspeed_mm = options.multimodal_inputs.map(|mm| match mm { - MultimodalData::TokenSpeed(data) => { - data.try_export_encoder_inputs_nixl_remote().into_proto() - } + MultimodalData::TokenSpeed(data) => data.into_proto(true), _ => 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 c137f9c3ab..ce1d49f046 100644 --- a/model_gateway/src/routers/grpc/epd_encode.rs +++ b/model_gateway/src/routers/grpc/epd_encode.rs @@ -15,11 +15,11 @@ use uuid::Uuid; use super::{ client::GrpcClient, context::{ClientSelection, WorkerSelection}, - multimodal::{assemble_tokenspeed_for_encode, MultimodalIntermediate}, + multimodal::{assemble_tokenspeed_for_encode, mm_rdma_exporter, MultimodalIntermediate}, proto_wrapper::{ cleanup_mm_shm_handles, cleanup_tokenspeed_items_encoder_shm, - collect_tokenspeed_multimodal_inputs_shm_handles, EncodeItemBootstrapInfo, - TokenSpeedMultimodalData, TokenSpeedMultimodalItem, + collect_tokenspeed_multimodal_inputs_shm_handles, stage_tokenspeed_tensor_rdma, + EncodeItemBootstrapInfo, TokenSpeedMultimodalData, TokenSpeedMultimodalItem, }, }; use crate::worker::DEFAULT_BOOTSTRAP_PORT; @@ -108,7 +108,12 @@ impl PreparedEncodeItem { .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); + // Stage with bootstrap_room as the slot key (load-bearing), so + // into_proto(false) below does not re-stage with a random key. + if let Some(exporter) = mm_rdma_exporter() { + item.encoder_input = + stage_tokenspeed_tensor_rdma(exporter, bootstrap_room, item.encoder_input); + } let request = tokenspeed_encoder::EncodeRequest { request_id: format!("encode-{}", Uuid::now_v7()), mm_inputs: Some( @@ -117,7 +122,7 @@ impl PreparedEncodeItem { shm_enabled: *shm_enabled, shm_min_bytes: *shm_min_bytes, } - .into_proto(), + .into_proto(false), ), items: vec![tokenspeed_encoder::EncodeItemAssignment { bootstrap_room }], }; diff --git a/model_gateway/src/routers/grpc/multimodal/assemble.rs b/model_gateway/src/routers/grpc/multimodal/assemble.rs index a18028392b..c05b6a1963 100644 --- a/model_gateway/src/routers/grpc/multimodal/assemble.rs +++ b/model_gateway/src/routers/grpc/multimodal/assemble.rs @@ -230,6 +230,8 @@ fn assemble_vllm( modality, shm_enabled: resolve_mm_shm_enabled(workers, false), shm_min_bytes: resolve_mm_shm_min_bytes(workers), + // vLLM workers cannot pull RDMA payloads yet. + rdma_enabled: false, }) } diff --git a/model_gateway/src/routers/grpc/proto_wrapper.rs b/model_gateway/src/routers/grpc/proto_wrapper.rs index c2d9472622..d8b4deb095 100644 --- a/model_gateway/src/routers/grpc/proto_wrapper.rs +++ b/model_gateway/src/routers/grpc/proto_wrapper.rs @@ -37,6 +37,7 @@ use smg_grpc_client::{ vllm_engine::AbortOnDropStream as VllmStream, vllm_proto::{self as vllm, generate_complete::MatchedStop as VllmMatchedStop}, }; +use smg_mm_rdma::RdmaExporter; use crate::routers::grpc::multimodal::mm_rdma_exporter; @@ -105,6 +106,9 @@ pub struct VllmMultimodalData { /// config or the environment. pub shm_enabled: bool, pub shm_min_bytes: usize, + /// Whether `pixel_values` may use the RDMA lane. Gated on the worker being able + /// to pull, so SMG never emits a `remote` payload a worker would reject. + pub rdma_enabled: bool, } /// TRT-LLM multimodal data: raw image bytes only. @@ -194,45 +198,6 @@ impl TokenSpeedTensor { } } - pub fn try_export_nixl_remote(self, slot_key: i64) -> Self { - let Some(exporter) = mm_rdma_exporter() else { - 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 exporter.export(slot_key, 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(), @@ -291,6 +256,15 @@ impl VllmMultimodalData { pub fn into_proto(self) -> vllm::MultimodalInputs { let shm_enabled = self.shm_enabled; let shm_min_bytes = self.shm_min_bytes; + // RDMA only for pixel_values (see TokenSpeed item conversion). + let pixel_rdma = if self.rdma_enabled { + mm_rdma_exporter().map(|exporter| MmRdmaExport { + exporter, + slot_key: rand::rng().random_range(0..i64::MAX), + }) + } else { + None + }; let model_specific_tensors = self .model_specific_tensors .into_iter() @@ -300,7 +274,12 @@ impl VllmMultimodalData { vllm::TensorData { shape: v.shape, dtype: v.dtype, - payload: Some(vllm_tensor_payload(v.data, shm_enabled, shm_min_bytes)), + payload: Some(vllm_tensor_payload( + v.data, + shm_enabled, + shm_min_bytes, + None, + )), }, ) }) @@ -320,6 +299,7 @@ impl VllmMultimodalData { self.pixel_values, shm_enabled, shm_min_bytes, + pixel_rdma, )), }), model_specific_tensors, @@ -344,55 +324,56 @@ 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 mm_rdma_exporter().is_none() { - return self; - } - for item in &mut self.items { - let slot_key = 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(slot_key); - } - 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 { + /// Convert to TokenSpeed proto MultimodalInputs. `rdma_enabled` stages each + /// inline encoder_input over RDMA with a random slot key (the general path); EPD + /// passes `false` since it stages explicitly with `bootstrap_room` first. + pub fn into_proto(self, rdma_enabled: bool) -> tokenspeed::MultimodalInputs { let shm_enabled = self.shm_enabled; let shm_min_bytes = self.shm_min_bytes; + let rdma_exporter = if rdma_enabled { + mm_rdma_exporter() + } else { + None + }; let items = self .items .into_iter() - .map(|item| item.into_proto(shm_enabled, shm_min_bytes)) + .map(|item| item.into_proto(shm_enabled, shm_min_bytes, rdma_exporter)) .collect(); tokenspeed::MultimodalInputs { items } } } impl TokenSpeedMultimodalItem { - fn into_proto(self, shm_enabled: bool, shm_min_bytes: usize) -> tokenspeed::MultimodalItem { + fn into_proto( + self, + shm_enabled: bool, + shm_min_bytes: usize, + rdma_exporter: Option<&RdmaExporter>, + ) -> tokenspeed::MultimodalItem { let placeholders = self .mm_placeholders .into_iter() .map(|(offset, length)| tokenspeed::PlaceholderRange { offset, length }) .collect::>(); + // No RDMA for the small model-specific side tensors: keeps the fixed slot + // pool for the pixel-heavy encoder_input. let model_specific_tensors = self .model_specific_tensors .into_iter() .map(|(k, v)| (k, tensor_bytes_to_tokenspeed(v, shm_enabled, shm_min_bytes))) .collect::>(); + let encoder_rdma = rdma_exporter.map(|exporter| MmRdmaExport { + exporter, + slot_key: rand::rng().random_range(0..i64::MAX), + }); let encoder_input = Some(tokenspeed_tensor_to_proto( self.encoder_input, shm_enabled, shm_min_bytes, + encoder_rdma, )); tokenspeed::MultimodalItem { @@ -414,6 +395,7 @@ fn tokenspeed_tensor_to_proto( value: TokenSpeedTensor, shm_enabled: bool, shm_min_bytes: usize, + rdma: Option>, ) -> tokenspeed::TensorData { use crate::observability::metrics::Metrics; let TokenSpeedTensor { @@ -422,9 +404,9 @@ fn tokenspeed_tensor_to_proto( dtype, } = value; let payload = match storage { - // Inline storage is metered inside tokenspeed_tensor_payload. + // Inline storage: RDMA/SHM/inline is resolved (and metered) here. TokenSpeedTensorStorage::Inline(data) => { - tokenspeed_tensor_payload(data, shm_enabled, shm_min_bytes) + tokenspeed_tensor_payload(data, shm_enabled, shm_min_bytes, rdma) } // Encoder input already written directly to SHM upstream — meter it here. TokenSpeedTensorStorage::Shm(handle) => { @@ -454,17 +436,74 @@ fn tensor_bytes_to_tokenspeed( tokenspeed::TensorData { shape, dtype, - payload: Some(tokenspeed_tensor_payload(data, shm_enabled, shm_min_bytes)), + payload: Some(tokenspeed_tensor_payload( + data, + shm_enabled, + shm_min_bytes, + None, + )), } } -/// Engine-neutral inline-vs-SHM decision for a multimodal tensor payload. The -/// raw bytes go inline unless SHM is enabled and the payload is at least -/// `min_bytes`; an SHM-write failure falls back to inline. `engine` labels the -/// metrics/logs. Each backend maps the result onto its own `TensorData` oneof. +/// Transport decision for one multimodal tensor; each backend maps it onto its +/// own `TensorData` oneof. enum MmTensorPayload { Inline(Vec), Shm(common::ShmHandle), + Remote(common::RemoteTensorHandle), +} + +/// A request to stage one tensor over RDMA. `slot_key` tags the descriptor: the +/// general path mints a random key, EPD uses the load-bearing `bootstrap_room`. +#[derive(Clone, Copy)] +struct MmRdmaExport<'a> { + exporter: &'a RdmaExporter, + slot_key: i64, +} + +/// Stage `data` into the RDMA arena under `slot_key`; `Ok(handle)`, or the bytes +/// back on failure (empty tensors are never staged). +fn export_rdma_tensor( + exporter: &RdmaExporter, + slot_key: i64, + data: Vec, +) -> Result> { + if data.is_empty() { + return Err(data); + } + let nbytes = data.len() as u64; + exporter + .export(slot_key, data) + .map(|descriptor| common::RemoteTensorHandle { + transport: "nixl".to_string(), + descriptor, + nbytes, + }) +} + +/// Stage an inline TokenSpeed encoder input over RDMA under an explicit `slot_key` +/// (EPD uses `bootstrap_room`). Non-inline tensors pass through untouched. +pub(crate) fn stage_tokenspeed_tensor_rdma( + exporter: &RdmaExporter, + slot_key: i64, + tensor: TokenSpeedTensor, +) -> TokenSpeedTensor { + let TokenSpeedTensor { + storage, + shape, + dtype, + } = tensor; + let TokenSpeedTensorStorage::Inline(data) = storage else { + return TokenSpeedTensor { + storage, + shape, + dtype, + }; + }; + match export_rdma_tensor(exporter, slot_key, data) { + Ok(handle) => TokenSpeedTensor::remote(handle, shape, dtype), + Err(data) => TokenSpeedTensor::inline(data, shape, dtype), + } } fn resolve_mm_tensor_payload( @@ -472,9 +511,27 @@ fn resolve_mm_tensor_payload( shm_enabled: bool, min_bytes: usize, engine: &'static str, + rdma: Option>, ) -> MmTensorPayload { use crate::observability::metrics::Metrics; let log_timing = log_tokenspeed_mm_timing_enabled(); + + // RDMA first; on export failure the bytes are handed back so we fall through to + // SHM/inline rather than drop the payload. + let data = match rdma { + Some(rdma) => { + let nbytes = data.len(); + match export_rdma_tensor(rdma.exporter, rdma.slot_key, data) { + Ok(handle) => { + Metrics::record_mm_tensor(engine, "remote", nbytes); + return MmTensorPayload::Remote(handle); + } + Err(data) => data, + } + } + None => data, + }; + let nbytes = data.len(); if !shm_enabled || nbytes < min_bytes { if log_timing { @@ -521,10 +578,12 @@ fn tokenspeed_tensor_payload( data: Vec, shm_enabled: bool, min_bytes: usize, + rdma: Option>, ) -> tokenspeed::tensor_data::Payload { - match resolve_mm_tensor_payload(data, shm_enabled, min_bytes, "tokenspeed") { + match resolve_mm_tensor_payload(data, shm_enabled, min_bytes, "tokenspeed", rdma) { MmTensorPayload::Inline(data) => tokenspeed::tensor_data::Payload::Inline(data), MmTensorPayload::Shm(handle) => tokenspeed::tensor_data::Payload::Shm(handle), + MmTensorPayload::Remote(handle) => tokenspeed::tensor_data::Payload::Remote(handle), } } @@ -532,10 +591,12 @@ fn vllm_tensor_payload( data: Vec, shm_enabled: bool, min_bytes: usize, + rdma: Option>, ) -> vllm::tensor_data::Payload { - match resolve_mm_tensor_payload(data, shm_enabled, min_bytes, "vllm") { + match resolve_mm_tensor_payload(data, shm_enabled, min_bytes, "vllm", rdma) { MmTensorPayload::Inline(data) => vllm::tensor_data::Payload::Inline(data), MmTensorPayload::Shm(handle) => vllm::tensor_data::Payload::Shm(handle), + MmTensorPayload::Remote(handle) => vllm::tensor_data::Payload::Remote(handle), } } @@ -2042,7 +2103,7 @@ mod tests { shm_enabled: false, shm_min_bytes: 0, } - .into_proto(); + .into_proto(false); assert_eq!(proto.items.len(), 1); let item = &proto.items[0]; @@ -2085,7 +2146,7 @@ mod tests { shm_enabled: false, shm_min_bytes: 0, } - .into_proto(); + .into_proto(false); assert_eq!(proto.items.len(), 1); let item = &proto.items[0]; @@ -2145,7 +2206,7 @@ mod tests { shm_enabled: true, shm_min_bytes: 0, } - .into_proto(); + .into_proto(false); let tensor = proto.items[0].encoder_input.as_ref().unwrap(); assert_eq!(tensor.shape, vec![1, 2]); @@ -2197,7 +2258,7 @@ mod tests { shm_enabled: true, shm_min_bytes: 0, } - .into_proto(); + .into_proto(false); let tensor = proto.items[0].encoder_input.as_ref().unwrap(); assert_eq!(tensor.shape, vec![1, 2]); @@ -2255,6 +2316,7 @@ mod tests { modality, shm_enabled: false, shm_min_bytes: 0, + rdma_enabled: false, } }