feat(multimodal): vLLM SHM tensor transport

crates/grpc_client/proto/vllm_engine.proto

@@ -95,11 +95,22 @@ message TokenizedInput {
   repeated uint32 input_ids = 2;  // Actual token IDs to process
 }
 
-// A typed tensor: raw little-endian bytes + shape + dtype.
+// A typed tensor descriptor. The raw bytes live in exactly one payload
+// transport; shape and dtype describe how the receiver interprets them.
 message TensorData {
-  bytes data = 1;              // Raw little-endian bytes (f32/i64/u32)
   repeated uint32 shape = 2;   // Dimension sizes
   string dtype = 3;            // "float32", "int64", "uint32"
+
+  oneof payload {
+    // Current path: raw little-endian bytes carried in the gRPC message
+    // (field 1, formerly `bytes data`, kept for wire compatibility).
+    bytes inline = 1;
+    // Same-host CPU shared memory path; preferred for large payloads when
+    // SMG and the vLLM worker share /dev/shm.
+    smg.grpc.common.ShmHandle shm = 4;
+    // Cross-node or non-shared-memory transport (e.g. NIXL). Not implemented.
+    smg.grpc.common.RemoteTensorHandle remote = 5;
+  }
 }
 
 message PlaceholderRange {
@@ -339,6 +350,11 @@ message GetServerInfoResponse {
   string kv_engine_id = 8;  // kv_transfer_config.engine_id, "" if not configured
 
   int32 data_parallel_size = 9;  // parallel_config.data_parallel_size (1 when DP is off)
+
+  // This worker's /dev/shm tmpfs identity (<boot_id>:<st_dev>), advertised so
+  // the router can verify a shared /dev/shm before using the SHM tensor
+  // transport under `auto`. Empty when it can't be determined.
+  string shm_namespace_id = 10;
 }
 
 // =====================

grpc_servicer/smg_grpc_servicer/mm_shm.py

@@ -0,0 +1,86 @@
+"""Shared multimodal SHM tensor-transport helpers (engine-neutral).
+
+Reads a ``TensorData`` payload that carries its raw little-endian bytes either
+inline in the gRPC message or via a same-host ``/dev/shm`` handle, and reports
+this process's ``/dev/shm`` namespace identity. Used by the vLLM and TokenSpeed
+gRPC servicers; their ``TensorData``/``ShmHandle`` messages share the ``payload``
+oneof shape (``inline`` | ``shm`` | ``remote``) via ``common.proto``.
+"""
+
+import os
+
+# Unlink each /dev/shm segment right after the worker reads it (default on) so
+# same-host SHM tensors don't accumulate. Disable with
+# TOKENSPEED_UNLINK_MM_SHM_AFTER_READ=0 (e.g. for debugging).
+UNLINK_MM_SHM_AFTER_READ = os.getenv("TOKENSPEED_UNLINK_MM_SHM_AFTER_READ", "1").lower() not in (
+    "0",
+    "false",
+    "no",
+)
+
+
+def tensor_payload_bytes(tensor_data) -> bytes:
+    """Raw bytes of a ``TensorData``, from whichever payload transport it carries."""
+    payload = tensor_data.WhichOneof("payload")
+    if payload == "inline":
+        return bytes(tensor_data.inline)
+    if payload == "shm":
+        return tensor_payload_bytes_from_shm(tensor_data.shm)
+    if payload == "remote":
+        raise ValueError("TensorData.remote payload is not implemented yet")
+    raise ValueError("TensorData payload is required")
+
+
+def tensor_payload_bytes_from_shm(shm_handle) -> bytes:
+    """Read ``nbytes`` from ``/dev/shm/<name>`` at ``offset`` (unlinking after read
+    when enabled)."""
+    name = validated_shm_name(shm_handle.name)
+
+    path = os.path.join("/dev/shm", name)
+    fd = None
+    try:
+        fd = os.open(path, os.O_RDONLY)
+        raw = os.pread(fd, int(shm_handle.nbytes), int(shm_handle.offset))
+    finally:
+        if fd is not None:
+            os.close(fd)
+        if fd is not None and UNLINK_MM_SHM_AFTER_READ:
+            try:
+                os.unlink(path)
+            except FileNotFoundError:
+                pass
+
+    if len(raw) != int(shm_handle.nbytes):
+        raise ValueError(
+            f"TensorData.shm byte length mismatch for name={shm_handle.name!r}: "
+            f"expected {int(shm_handle.nbytes)}, got {len(raw)}"
+        )
+    return raw
+
+
+def validated_shm_name(name: str) -> str:
+    """Reject path-traversal / absolute / empty SHM names before opening."""
+    name = name.lstrip("/")
+    if not name or "/" in name or name in (".", "..") or "\x00" in name:
+        raise ValueError(f"Invalid TensorData.shm name: {name!r}")
+    return name
+
+
+def shm_namespace_id() -> str:
+    """Identity of this process's ``/dev/shm`` tmpfs: ``<boot_id>:<st_dev>``.
+
+    ``boot_id`` (``/proc/sys/kernel/random/boot_id``) is not namespaced, so it
+    pins the host; ``st_dev`` is the tmpfs superblock device backing
+    ``/dev/shm``. Two processes share ``/dev/shm`` iff both match -- including
+    separate containers sharing it via ``--ipc``/bind-mount, where mount
+    namespaces differ but the underlying superblock (``st_dev``) is the same. The
+    router compares this token to its own to decide the SHM tensor transport.
+    Empty string if it can't be determined.
+    """
+    try:
+        with open("/proc/sys/kernel/random/boot_id", encoding="ascii") as f:
+            boot_id = f.read().strip()
+        shm_dev = os.stat("/dev/shm").st_dev
+        return f"{boot_id}:{shm_dev}"
+    except OSError:
+        return ""

grpc_servicer/smg_grpc_servicer/vllm/servicer.py

@@ -21,6 +21,7 @@
 import zmq.asyncio
 from smg_grpc_proto import vllm_engine_pb2, vllm_engine_pb2_grpc
 from smg_grpc_proto.generated import common_pb2
+from smg_grpc_servicer import mm_shm
 from transformers import BatchFeature
 from vllm import PoolingParams, SamplingParams, TokensPrompt
 from vllm.distributed.kv_events import KVEventBatch
@@ -78,7 +79,8 @@ def _tensor_from_proto(td: vllm_engine_pb2.TensorData) -> torch.Tensor:
     torch_dtype = _PROTO_DTYPE_MAP.get(td.dtype)
     if torch_dtype is None:
         raise ValueError(f"Unsupported proto tensor dtype: {td.dtype!r}")
-    return torch.frombuffer(bytearray(td.data), dtype=torch_dtype).reshape(*td.shape)
+    payload = mm_shm.tensor_payload_bytes(td)
+    return torch.frombuffer(bytearray(payload), dtype=torch_dtype).reshape(*td.shape)
 
 
 try:
@@ -487,6 +489,7 @@ async def GetServerInfo(
             kv_role=kv_role,
             kv_engine_id=kv_engine_id,
             data_parallel_size=parallel.data_parallel_size,
+            shm_namespace_id=mm_shm.shm_namespace_id(),
         )
 
     async def GetLoads(

model_gateway/src/routers/grpc/client.rs

@@ -14,9 +14,9 @@ use smg_grpc_client::{
 
 use crate::routers::grpc::{
     proto_wrapper::{
-        cleanup_tokenspeed_shm_handles, collect_tokenspeed_generate_request_shm_handles,
-        finish_tokenspeed_request, ProtoEmbedComplete, ProtoEmbedRequest, ProtoGenerateRequest,
-        ProtoStream,
+        cleanup_mm_shm_handles, collect_tokenspeed_generate_request_shm_handles,
+        collect_vllm_generate_request_shm_handles, finish_tokenspeed_request, ProtoEmbedComplete,
+        ProtoEmbedRequest, ProtoGenerateRequest, ProtoStream,
     },
     MultimodalData,
 };
@@ -397,8 +397,14 @@ impl GrpcClient {
                 Ok(ProtoStream::Sglang(stream))
             }
             (Self::Vllm(client), ProtoGenerateRequest::Vllm(boxed_req)) => {
-                let stream = client.generate(*boxed_req).await?;
-                Ok(ProtoStream::Vllm(stream))
+                let shm_handles = collect_vllm_generate_request_shm_handles(&boxed_req);
+                match client.generate(*boxed_req).await {
+                    Ok(stream) => Ok(ProtoStream::Vllm(stream)),
+                    Err(error) => {
+                        cleanup_mm_shm_handles(&shm_handles);
+                        Err(error)
+                    }
+                }
             }
             (Self::Trtllm(client), ProtoGenerateRequest::Trtllm(boxed_req)) => {
                 let stream = client.generate(*boxed_req).await?;
@@ -413,7 +419,7 @@ impl GrpcClient {
                 match client.generate(*boxed_req).await {
                     Ok(stream) => Ok(ProtoStream::TokenSpeed(stream)),
                     Err(error) => {
-                        cleanup_tokenspeed_shm_handles(&shm_handles);
+                        cleanup_mm_shm_handles(&shm_handles);
                         Err(error)
                     }
                 }

model_gateway/src/routers/grpc/epd_encode.rs

@@ -17,7 +17,7 @@ use super::{
     context::{ClientSelection, WorkerSelection},
     multimodal::{assemble_tokenspeed, MultimodalIntermediate, PrecomputedMultimodalIntermediate},
     proto_wrapper::{
-        cleanup_tokenspeed_items_encoder_shm, cleanup_tokenspeed_shm_handles,
+        cleanup_mm_shm_handles, cleanup_tokenspeed_items_encoder_shm,
         collect_tokenspeed_multimodal_inputs_shm_handles, EncodeItemBootstrapInfo,
         TokenSpeedMultimodalData, TokenSpeedMultimodalItem,
     },
@@ -136,7 +136,7 @@ struct TokenSpeedShmCleanupGuard(Vec<common_proto::ShmHandle>);
 
 impl Drop for TokenSpeedShmCleanupGuard {
     fn drop(&mut self) {
-        cleanup_tokenspeed_shm_handles(&self.0);
+        cleanup_mm_shm_handles(&self.0);
     }
 }
 

model_gateway/src/routers/grpc/multimodal/assemble.rs

@@ -72,7 +72,7 @@ async fn assemble_multimodal_data_impl(
             }
             GrpcClient::Vllm(_) => {
                 ensure_image_only(&precomputed, "vLLM")?;
-                Ok(MultimodalData::Vllm(assemble_vllm(precomputed)))
+                Ok(MultimodalData::Vllm(assemble_vllm(precomputed, workers)))
             }
             GrpcClient::Trtllm(_) => {
                 ensure_image_only(&precomputed, "TRT-LLM")?;
@@ -168,7 +168,10 @@ fn assemble_sglang(intermediate: PrecomputedMultimodalIntermediate) -> SglangMul
     }
 }
 
-fn assemble_vllm(intermediate: PrecomputedMultimodalIntermediate) -> VllmMultimodalData {
+fn assemble_vllm(
+    intermediate: PrecomputedMultimodalIntermediate,
+    workers: Option<&WorkerSelection>,
+) -> VllmMultimodalData {
     let (pixel_values, pixel_values_shape) = serialize_encoder_input(&intermediate.preprocessed);
     let model_specific_tensors = serialize_model_specific(intermediate.preprocessed.model_specific);
     let mm_hashes = intermediate.images.iter().map(|f| f.hash.clone()).collect();
@@ -190,6 +193,8 @@ fn assemble_vllm(intermediate: PrecomputedMultimodalIntermediate) -> VllmMultimo
         batched_keys,
         flat_keys,
         keep_on_cpu_keys: intermediate.keep_on_cpu_keys,
+        shm_enabled: resolve_mm_shm_enabled(workers, false),
+        shm_min_bytes: resolve_mm_shm_min_bytes(workers),
     }
 }