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,100 @@
+"""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
+import stat
+
+# 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:
+        # O_NOFOLLOW: never follow a symlink at the final path component, so a
+        # crafted name that resolves to a pre-existing symlink in /dev/shm can't
+        # redirect the read. Then require a regular file (fail closed otherwise).
+        fd = os.open(path, os.O_RDONLY | os.O_NOFOLLOW)
+        if not stat.S_ISREG(os.fstat(fd).st_mode):
+            raise ValueError(f"TensorData.shm is not a regular file: {shm_handle.name!r}")
+        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
+
+
+_shm_namespace_id_cache: str | None = None
+
+
+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. Cached: both components are static
+    for the process lifetime, and this is read on every GetServerInfo.
+    """
+    global _shm_namespace_id_cache
+    if _shm_namespace_id_cache is not None:
+        return _shm_namespace_id_cache
+    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
+        _shm_namespace_id_cache = f"{boot_id}:{shm_dev}"
+    except OSError:
+        _shm_namespace_id_cache = ""
+    return _shm_namespace_id_cache

grpc_servicer/smg_grpc_servicer/vllm/servicer.py

@@ -37,6 +37,7 @@
 from vllm.outputs import CompletionOutput, RequestOutput
 from vllm.sampling_params import RequestOutputKind, StructuredOutputsParams
 
+from smg_grpc_servicer import mm_shm
 from smg_grpc_servicer.tokenizer_bundle import CHUNK_SIZE, build_tokenizer_zip
 from smg_grpc_servicer.vllm.kv_events import (
     endpoint_for_rank,
@@ -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, finish_vllm_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)
                     }
                 }
@@ -489,15 +495,16 @@ impl GrpcClient {
                     MultimodalData::Vllm(data) => data.into_proto(),
                     _ => unreachable!("caller guarantees matching variant"),
                 });
-                let req = client.build_generate_request_from_chat(
-                    request_id,
-                    body,
-                    processed_text,
-                    token_ids,
-                    vllm_mm,
-                    options.tool_constraints,
-                )?;
-                Ok(ProtoGenerateRequest::Vllm(Box::new(req)))
+                finish_vllm_request(vllm_mm, |mm| {
+                    client.build_generate_request_from_chat(
+                        request_id,
+                        body,
+                        processed_text,
+                        token_ids,
+                        mm,
+                        options.tool_constraints,
+                    )
+                })
             }
             Self::Trtllm(client) => {
                 let trtllm_mm = options.multimodal_inputs.map(|mm| match mm {
@@ -580,15 +587,16 @@ impl GrpcClient {
                     MultimodalData::Vllm(data) => data.into_proto(),
                     _ => unreachable!("caller guarantees matching variant"),
                 });
-                let req = client.build_generate_request_from_messages(
-                    request_id,
-                    body,
-                    processed_text,
-                    token_ids,
-                    vllm_mm,
-                    options.tool_constraints,
-                )?;
-                Ok(ProtoGenerateRequest::Vllm(Box::new(req)))
+                finish_vllm_request(vllm_mm, |mm| {
+                    client.build_generate_request_from_messages(
+                        request_id,
+                        body,
+                        processed_text,
+                        token_ids,
+                        mm,
+                        options.tool_constraints,
+                    )
+                })
             }
             Self::Trtllm(client) => {
                 let trtllm_mm = options.multimodal_inputs.map(|mm| match mm {

model_gateway/src/routers/grpc/common/stages/request_execution.rs

@@ -657,6 +657,11 @@ impl RequestExecutionStage {
         // Decode reuses proto_request as-is; same request_id as the prefill leg is
         // load-bearing for NIXL P/D correlation on vLLM < 0.13
         let mut decode_request = proto_request;
+        // Decode doesn't run the vision encoder (it receives KV via the P/D
+        // transfer), so drop the multimodal inputs — mirrors the parallel PD
+        // path. Load-bearing for SHM: prefill already read and unlinked the
+        // /dev/shm segments, so a reused ShmHandle here would be unreadable.
+        decode_request.clear_mm_pixel_values();
         if let Some(rank) = workers.decode_worker().and_then(|w| w.dp_rank()) {
             decode_request.set_data_parallel_rank(rank as i32);
         }

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),
     }
 }