feat(multimodal): vLLM SHM tensor transport - #1893
Conversation
Bring the same-host /dev/shm tensor transport to the vLLM backend, reusing the
engine-neutral ShmHandle (common.proto) and the configurable transport layer.
Proto (vllm_engine.proto):
- TensorData: `bytes data = 1` -> `oneof payload { bytes inline = 1;
smg.grpc.common.ShmHandle shm = 4; smg.grpc.common.RemoteTensorHandle
remote = 5; }`. Field 1 stays inline for wire compatibility.
- GetServerInfoResponse: add `shm_namespace_id = 10` so the router can verify a
shared /dev/shm under `auto`.
Gateway (Rust):
- Extract an engine-neutral `resolve_mm_tensor_payload` (inline-vs-SHM decision +
metrics), used by both `tokenspeed_tensor_payload` and the new
`vllm_tensor_payload`.
- `VllmMultimodalData` carries the resolved shm_enabled / shm_min_bytes;
`assemble_vllm` resolves them from the transport config + worker; `into_proto`
emits the oneof payload.
- Add `collect_vllm_*_shm_handles`; the vLLM client cleans up SHM handles on
send failure (mirrors TokenSpeed). Rename the engine-neutral
`cleanup_tokenspeed_shm_handles` -> `cleanup_mm_shm_handles`.
Servicer (Python):
- New shared `mm_shm` module (reads a TensorData payload inline-or-SHM, computes
the /dev/shm namespace id). The vLLM servicer reads tensors through it and
advertises `shm_namespace_id` in GetServerInfo.
Default transport is `inline`, so existing deployments are unaffected. Migrating
the TokenSpeed servicer onto the shared `mm_shm` module (dedup) is a follow-up.
e2e for the SHM path is deferred to the multimodal coverage PR.
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
github-actions
Bot
added
grpc
📝 WalkthroughWalkthroughThis PR adds SHM-aware tensor payload transport to the vLLM proto, wires Python tensor decoding and server info reporting through shared-memory helpers, and updates the Rust gateway to build, transport, clean up, and reuse multimodal SHM payloads across vLLM request paths. ChangesMultimodal SHM Tensor Transport
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Client as model_gateway
participant Wrapper as proto_wrapper.rs
participant Servicer as vllm/servicer.py
participant Shm as /dev/shm
Client->>Wrapper: resolve_mm_tensor_payload(...)
Wrapper->>Shm: write or reference SHM payload
Wrapper-->>Client: ProtoGenerateRequest
Client->>Servicer: gRPC request with TensorData
Servicer->>Servicer: mm_shm.tensor_payload_bytes(td)
Servicer->>Shm: os.pread(fd, nbytes, offset)
Shm-->>Servicer: raw tensor bytes
Client->>Wrapper: cleanup_mm_shm_handles(handles) on failure
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements shared memory (SHM) tensor-transport support for vLLM multimodal inputs, matching the existing TokenSpeed SHM transport mechanism. It updates protobuf definitions, introduces a shared Python helper (mm_shm.py) for reading payloads and identifying /dev/shm namespaces, and refactors the Rust model gateway to handle vLLM SHM transport and cleanup. The review feedback highlights several performance optimization opportunities: caching the SHM namespace ID in Python to avoid repeated system calls, passing the payload directly to torch.frombuffer to achieve true zero-copy deserialization, and caching the environment variable lookup in Rust using OnceLock to prevent lock contention in hot paths.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| payload = mm_shm.tensor_payload_bytes(td) | ||
| return torch.frombuffer(bytearray(payload), dtype=torch_dtype).reshape(*td.shape) |
There was a problem hiding this comment.
Using bytearray(payload) creates a mutable copy of the entire tensor data in CPU memory, which introduces an unnecessary copy and defeats the zero-copy benefit of the shared memory (SHM) transport. Since torch.frombuffer natively supports read-only buffers (like bytes) and returns a read-only tensor, you can pass payload directly to torch.frombuffer to achieve true zero-copy deserialization.
| payload = mm_shm.tensor_payload_bytes(td) | |
| return torch.frombuffer(bytearray(payload), dtype=torch_dtype).reshape(*td.shape) | |
| payload = mm_shm.tensor_payload_bytes(td) | |
| return torch.frombuffer(payload, dtype=torch_dtype).reshape(*td.shape) |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 6a59bde072
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| batched_keys, | ||
| flat_keys, | ||
| keep_on_cpu_keys: intermediate.keep_on_cpu_keys, | ||
| shm_enabled: resolve_mm_shm_enabled(workers, false), |
There was a problem hiding this comment.
In vLLM PD, execute_sequential_pd sends a cloned request to prefill and later reuses the original request as the decode request, so both legs carry the same mm_inputs. With SHM enabled here whenever the selected workers share /dev/shm, the prefill servicer reads and unlinks each segment, then the decode servicer receives the same ShmHandle and fails to open it before generation. Please either disable SHM for the multi-consumer vLLM PD path, create separate SHM payloads per leg, or keep the segment alive until both legs have consumed it.
Useful? React with 👍 / 👎.
| data: self.pixel_values, | ||
| shape: self.pixel_values_shape, | ||
| dtype: "float32".to_string(), | ||
| payload: Some(vllm_tensor_payload( |
There was a problem hiding this comment.
When vllm_tensor_payload chooses SHM, this call creates /dev/shm files during into_proto(), before the vLLM chat/messages builders run sampling and tool-constraint validation. Those builders still return with ? and do not use a finish_* cleanup wrapper like TokenSpeed, so invalid multimodal requests with SHM enabled leak one file per tensor before the request is ever sent. Please collect and unlink the vLLM handles on build errors as well as on send errors.
Useful? React with 👍 / 👎.
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
| (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); |
There was a problem hiding this comment.
🔴 Important: SHM cleanup covers the send-failure path here, but the build-failure path is unguarded. When into_proto() writes SHM segments during build_generate_request (e.g. at client.rs:495), then build_grpc_sampling_params_from_chat fails via ?, the SHM files leak — the segments are dropped without cleanup and generate() never runs.
TokenSpeed handles this with finish_tokenspeed_request (proto_wrapper.rs:675), which collects SHM handles before the build and cleans up on Err. The vLLM path needs an equivalent guard — either a finish_vllm_request wrapper or an explicit collect-and-cleanup around the into_proto() + build sequence in each build_generate_request arm.
| pub fn cleanup_mm_shm_handles(handles: &[common::ShmHandle]) { | ||
| for handle in handles { | ||
| let Some(name) = validate_tokenspeed_shm_name_for_cleanup(&handle.name) else { | ||
| tracing::warn!( |
There was a problem hiding this comment.
🟡 Nit: cleanup_mm_shm_handles was renamed to be engine-neutral, but its internals still reference "TokenSpeed" — the validation function (validate_tokenspeed_shm_name_for_cleanup) and the log messages ("invalid TokenSpeed SHM name", "Failed to cleanup TokenSpeed SHM file"). When debugging vLLM SHM cleanup issues, these messages will be confusing.
Since the naming prefix (smg-tokenspeed-) is shared for now (both engines use write_tokenspeed_shm), the behavior is correct — this is purely about log/diagnostic clarity. Consider updating the messages to say "multimodal SHM" or parameterizing by engine, similar to what resolve_mm_tensor_payload does.
| # 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 ( |
There was a problem hiding this comment.
🟡 Nit: The env var is TOKENSPEED_UNLINK_MM_SHM_AFTER_READ but this module is engine-neutral (used by both vLLM and TokenSpeed servicers). A SMG_UNLINK_MM_SHM_AFTER_READ name would be consistent with the module's scope. Not urgent — fine to defer if you plan to rename during the TokenSpeed servicer migration (mentioned in follow-ups).
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@grpc_servicer/smg_grpc_servicer/mm_shm.py`:
- Line 10: The `/dev/shm` file open path is still vulnerable to symlink
traversal even after `validated_shm_name()`. Update the shm-open logic in
`mm_shm.py` to use `os.open(..., os.O_RDONLY | os.O_NOFOLLOW)` and then verify
the opened descriptor is a regular file before reading. If the target is a
symlink or any non-regular file, reject it and fail closed.
In `@model_gateway/src/routers/grpc/proto_wrapper.rs`:
- Around line 237-248: Add cleanup for SHM-backed payloads in the vLLM request
build path, since `into_proto` in `GrpcProtoWrapper` can create `/dev/shm` files
before request construction succeeds. Update the vLLM chat/messages builders to
wrap the request creation logic in a helper similar to
`finish_tokenspeed_request`, and ensure any `Err` from `into_proto` or later
request assembly triggers cleanup of the SHM handles so leaked files are removed
when the Python worker never receives the request.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 43a8cdd1-3a40-48ae-b5a7-6ce8effa7298
📒 Files selected for processing (7)
crates/grpc_client/proto/vllm_engine.protogrpc_servicer/smg_grpc_servicer/mm_shm.pygrpc_servicer/smg_grpc_servicer/vllm/servicer.pymodel_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/epd_encode.rsmodel_gateway/src/routers/grpc/multimodal/assemble.rsmodel_gateway/src/routers/grpc/proto_wrapper.rs
| let shm_enabled = self.shm_enabled; | ||
| let shm_min_bytes = self.shm_min_bytes; | ||
| let model_specific_tensors = self | ||
| .model_specific_tensors | ||
| .into_iter() | ||
| .map(|(k, v)| { | ||
| ( | ||
| k, | ||
| vllm::TensorData { | ||
| data: v.data, | ||
| shape: v.shape, | ||
| dtype: v.dtype, | ||
| payload: Some(vllm_tensor_payload(v.data, shm_enabled, shm_min_bytes)), |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Add vLLM build-failure cleanup for SHM handles.
Line 248 and Line 264 can materialize /dev/shm files during into_proto. If vLLM request construction later returns Err, the request is never sent for the Python worker to read/unlink it; unlike TokenSpeed, the shown vLLM build paths do not wrap request construction with cleanup.
Proposed direction
+pub(crate) fn finish_vllm_request(
+ vllm_mm: Option<vllm::MultimodalInputs>,
+ build: impl FnOnce(Option<vllm::MultimodalInputs>) -> Result<vllm::GenerateRequest, String>,
+) -> Result<ProtoGenerateRequest, String> {
+ let shm_handles = vllm_mm
+ .as_ref()
+ .map(collect_vllm_multimodal_inputs_shm_handles)
+ .unwrap_or_default();
+ match build(vllm_mm) {
+ Ok(req) => Ok(ProtoGenerateRequest::Vllm(Box::new(req))),
+ Err(error) => {
+ cleanup_mm_shm_handles(&shm_handles);
+ Err(error)
+ }
+ }
+}Then use this wrapper in the vLLM chat/messages request builders, mirroring finish_tokenspeed_request.
Also applies to: 260-268
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@model_gateway/src/routers/grpc/proto_wrapper.rs` around lines 237 - 248, Add
cleanup for SHM-backed payloads in the vLLM request build path, since
`into_proto` in `GrpcProtoWrapper` can create `/dev/shm` files before request
construction succeeds. Update the vLLM chat/messages builders to wrap the
request creation logic in a helper similar to `finish_tokenspeed_request`, and
ensure any `Err` from `into_proto` or later request assembly triggers cleanup of
the SHM handles so leaked files are removed when the Python worker never
receives the request.
- Sequential vLLM PD: clear the decode leg's multimodal inputs (mirrors the parallel PD path). The decode leg doesn't run the vision encoder, and with SHM the prefill leg already read+unlinked the /dev/shm segments, so a reused ShmHandle on decode would be unreadable. (codex P1) - vLLM build path: add `finish_vllm_request`, which collects the SHM handles before building the request and unlinks them if the build fails (sampling / tool-constraint validation). `into_proto` can write /dev/shm files before the request is assembled, so a build error would otherwise leak them. Mirrors `finish_tokenspeed_request`. (claude, coderabbit, codex) - mm_shm read: open with O_NOFOLLOW and require a regular file, so a crafted TensorData.shm name resolving to a pre-existing /dev/shm symlink can't redirect the read. (coderabbit) - mm_shm: cache the /dev/shm namespace id (static per process, read on every GetServerInfo). (gemini) - proto_wrapper: cache SMG_LOG_MM_TIMING via OnceLock instead of re-reading the env per tensor payload. (gemini) - Fix ruff import ordering (mm_shm is first-party). Deferred to the TokenSpeed-servicer migration follow-up: rename the shared cleanup's TokenSpeed-flavored log/validation strings and the TOKENSPEED_UNLINK_MM_SHM_AFTER_READ env var to engine-neutral names. Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
|
Thanks all — addressed in 0a7d4ad:
Declined / deferred:
Verified: |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 0a7d4ade9c
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
| // 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(); |
There was a problem hiding this comment.
For vLLM PD requests with sampling n > 1, relay_kv_params is false, so the decode leg is intentionally not handed prefill KV and the NIXL path logs that decode will recompute the prompt locally. Clearing mm_inputs here means a multimodal decode request reaches vLLM with the expanded image tokens but without the image tensors/model-specific tensors, so NIXL n > 1 image requests will fail or generate from an incomplete prompt. Only drop these tensors on paths where decode actually consumes transferred KV, or give the recompute path its own payloads.
Useful? React with 👍 / 👎.
| 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(), |
There was a problem hiding this comment.
This servicer now constructs GetServerInfoResponse with shm_namespace_id, and _tensor_from_proto also calls WhichOneof("payload"), but grpc_servicer/pyproject.toml still allows smg-grpc-proto>=0.4.11 while the proto package version in this repo was not bumped past the pre-change 0.4.12. A normal package install can therefore pair the new servicer with generated stubs that do not define these fields, causing worker registration to raise TypeError: Protocol message GetServerInfoResponse has no "shm_namespace_id" field before SHM is even negotiated. Please bump/pin smg-grpc-proto to a new release before using the new field.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@grpc_servicer/smg_grpc_servicer/mm_shm.py`:
- Around line 46-48: Avoid blocking when opening shared-memory paths that may
point to FIFOs. In mm_shm.py, update the os.open call in the TensorData.shm path
handling to include O_NONBLOCK alongside O_RDONLY and O_NOFOLLOW, so the
subsequent stat.S_ISREG check can safely reject non-regular files without
hanging. Keep the fix localized to the same open-and-validate logic that uses
shm_handle.name and os.fstat(fd).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 2a8e02c9-f280-46fe-80db-ce83ec109637
📒 Files selected for processing (5)
grpc_servicer/smg_grpc_servicer/mm_shm.pygrpc_servicer/smg_grpc_servicer/vllm/servicer.pymodel_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/common/stages/request_execution.rsmodel_gateway/src/routers/grpc/proto_wrapper.rs
| 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}") |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
sed -n '1,220p' grpc_servicer/smg_grpc_servicer/mm_shm.pyRepository: lightseekorg/smg
Length of output: 1955
🏁 Script executed:
python3 - <<'PY'
import os, tempfile, multiprocessing, time, errno, stat, sys
td = tempfile.mkdtemp(prefix="fifo-probe-")
fifo = os.path.join(td, "x")
reg = os.path.join(td, "r")
os.mkfifo(fifo, 0o600)
with open(reg, "wb") as f:
f.write(b"hi")
def try_open(path, flags, q):
try:
fd = os.open(path, flags)
try:
mode = stat.S_IFMT(os.fstat(fd).st_mode)
q.put(("ok", mode))
finally:
os.close(fd)
except Exception as e:
q.put(("err", type(e).__name__, getattr(e, "errno", None), str(e)))
def run(path, flags, timeout=1.0):
q = multiprocessing.Queue()
p = multiprocessing.Process(target=try_open, args=(path, flags, q))
p.start()
p.join(timeout)
if p.is_alive():
p.terminate()
p.join()
return ("timeout",)
return q.get() if not q.empty() else ("no-result",)
cases = [
("fifo_rd", fifo, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)),
("fifo_rd_nonblock", fifo, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | os.O_NONBLOCK),
("reg_rd", reg, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0)),
("reg_rd_nonblock", reg, os.O_RDONLY | getattr(os, "O_NOFOLLOW", 0) | os.O_NONBLOCK),
]
for name, path, flags in cases:
print(name, run(path, flags))
PYRepository: lightseekorg/smg
Length of output: 1955
Avoid blocking FIFO opens
os.open(path, os.O_RDONLY | os.O_NOFOLLOW) can still block if /dev/shm contains a FIFO at path. Add O_NONBLOCK so the regular-file check can reject non-regular targets without hanging.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@grpc_servicer/smg_grpc_servicer/mm_shm.py` around lines 46 - 48, Avoid
blocking when opening shared-memory paths that may point to FIFOs. In mm_shm.py,
update the os.open call in the TensorData.shm path handling to include
O_NONBLOCK alongside O_RDONLY and O_NOFOLLOW, so the subsequent stat.S_ISREG
check can safely reject non-regular files without hanging. Keep the fix
localized to the same open-and-validate logic that uses shm_handle.name and
os.fstat(fd).
Enable video inputs on the vLLM gRPC path, previously rejected by ensure_image_only. Video pixel tensors ride the same inline/SHM transport as images (#1893), routed to vLLM's video modality. Instead of a one-off is_video bool, hoist the Modality enum into common.proto (like ShmHandle in #1) so the single-modality, precomputed-tensor engines (vLLM + TokenSpeed) share one modality type; TokenSpeed's proto now references smg.grpc.common.Modality. (SGLang keeps its string `modalities` for mixed-modality inputs; converging it onto the common enum is a follow-up.) - proto: Modality enum -> common.proto; tokenspeed references it; vLLM MultimodalInputs gets `common.Modality modality = 10` - assemble: vLLM accepts image or video; assemble_vllm maps the modality and takes mm_hashes from videos vs images. Mixed rejected upstream in process. - servicer: video -> pixel_values_videos under vLLM's `video` MultiModalFieldConfig; tokenspeed servicer reads modality from common_pb2 - test: modality proto round-trip Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
Enable video inputs on the vLLM gRPC path, previously rejected by ensure_image_only. Video pixel tensors ride the same inline/SHM transport as images (#1893), routed to vLLM's video modality. Instead of a one-off is_video bool, hoist the Modality enum into common.proto (like ShmHandle in #1) so the single-modality, precomputed-tensor engines (vLLM + TokenSpeed) share one modality type; TokenSpeed's proto now references smg.grpc.common.Modality. (SGLang keeps its string `modalities` for mixed-modality inputs; converging it onto the common enum is a follow-up.) - proto: Modality enum -> common.proto; tokenspeed references it; vLLM MultimodalInputs gets `common.Modality modality = 10` - assemble: vLLM accepts image or video; assemble_vllm maps the modality and takes mm_hashes from videos vs images. Mixed rejected upstream in process. - servicer: video -> pixel_values_videos under vLLM's `video` MultiModalFieldConfig; tokenspeed servicer reads modality from common_pb2 - test: modality proto round-trip Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
1 participant
Description
Problem
The
/dev/shmtensor transport (added for TokenSpeed) and the engine-neutralShmHandle(hoisted tocommon.proto) + configurable transport layer are allin place, but vLLM still only accepts multimodal tensors inline in the gRPC
message. Large image tensors to a co-located vLLM worker pay a full gRPC copy.
Solution
Bring the same-host SHM transport to vLLM, reusing the common
ShmHandleandthe
--multimodal-tensor-transportconfig. Default isinline, so existingdeployments are unaffected.
Changes
Proto (
vllm_engine.proto, already importscommon.proto):TensorData:bytes data = 1→oneof payload { bytes inline = 1; smg.grpc.common.ShmHandle shm = 4; smg.grpc.common.RemoteTensorHandle remote = 5; }. Field 1 staysinlinefor wire compatibility.GetServerInfoResponse:+ shm_namespace_id = 10(so the router can verify ashared
/dev/shmunderauto).Gateway (Rust):
resolve_mm_tensor_payload(the inline-vs-SHMdecision + metrics); both
tokenspeed_tensor_payloadand the newvllm_tensor_payloadmap its result onto their own oneof.VllmMultimodalDatacarries resolvedshm_enabled/shm_min_bytes;assemble_vllmresolves them (transport config → worker override);into_protoemits the oneof payload.
collect_vllm_*_shm_handles+ the vLLM client cleans up SHM on send failure(mirrors TokenSpeed). Renamed the engine-neutral
cleanup_tokenspeed_shm_handles→
cleanup_mm_shm_handles.Servicer (Python):
mm_shmmodule (reads aTensorDatapayload inline-or-SHM,validates SHM names, unlinks after read, computes the
/dev/shmnamespace id).mm_shm(theoneofrenames.data→
.inline, so this is required) and advertisesshm_namespace_id.Test Plan
cargo clippy -p smg --all-targets -- -D warnings— cleancargo test -p smg --lib— 1147 pass (one unrelatedmiddleware::metricsinterner test is a known parallel-test flake; passes in isolation)
cargo +nightly fmt— cleanvllm_engine_pb2.TensorDataexposes the
inline|shm|remoteoneof (shm→common_pb2.ShmHandle) andGetServerInfoResponse.shm_namespace_idexistspython -m py_compileonmm_shm.py+ the vLLM servicer — OKe2e for the SHM path (co-located worker, assert SHM is taken) is deferred to the
multimodal coverage PR.
Follow-ups
mm_shmmodule (mechanicaldedup; kept out of this PR to avoid churning the critical servicer).
is_video→pixel_values_videos) — PR #3b.Checklist
cargo +nightly fmt+clippy -- -D warningscleaninlinekeeps field 1)/dev/shmleak)Summary by CodeRabbit