perf(tokenspeed): optimize multimodal tensor transport - #1604
Merged
slin1237 merged 28 commits intoJun 22, 2026
Conversation
|
Caution Review failedPull request was closed or merged during review Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds SHM-backed tensor transport for TokenSpeed (gateway write + servicer read), PIL-exact bicubic resize and libjpeg-turbo decode for Qwen-VL preprocessing, media-before-text chat reordering, and comprehensive timing/metrics instrumentation across all transport and preprocessing paths. ChangesTokenSpeed SHM Transport, PIL Preprocessing, and Media Ordering
Sequence Diagram(s)sequenceDiagram
participant Client as GrpcClient
participant MW as multimodal.rs
participant PW as proto_wrapper.rs
participant SHM as /dev/shm
participant TS as TokenSpeed Servicer
Client->>MW: assemble_tokenspeed(items, worker_url)
MW->>PW: resolve_tokenspeed_shm_enabled(worker_url)
PW-->>MW: shm_enabled
loop per encoder item
MW->>PW: serialize_array_as_tokenspeed_tensor(array, shm_enabled)
alt shm_enabled and bytes >= threshold
PW->>SHM: write_tokenspeed_shm_with(nbytes, write_fn)
SHM-->>PW: ShmHandle or error
alt write succeeded
PW->>MW: record_mm_tensor(shm)
PW-->>MW: TokenSpeedTensor::Shm
else write failed
PW->>MW: record_mm_shm_write_failure()
PW-->>MW: TokenSpeedTensor::Inline fallback
end
else skip SHM attempt
PW->>MW: record_mm_tensor(inline)
PW-->>MW: TokenSpeedTensor::Inline
end
end
MW-->>Client: TokenSpeedMultimodalData{shm_enabled}
Client->>TS: generate RPC with Payload::Shm or Payload::Inline
alt Payload::Shm
TS->>SHM: os.pread(shm_name, nbytes)
SHM-->>TS: tensor bytes
TS->>TS: validate byte length
TS->>SHM: optional unlink(shm_name)
else Payload::Inline
TS->>TS: use inline bytes directly
end
TS-->>Client: stream response
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Possibly related issues
Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
github-actions
Bot
added
python-bindings
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces comprehensive video and multimodal processing support across the repository. Key changes include updating the gRPC protobuf definitions to support itemized multimodal payloads (images, audio, and video) and shared memory (SHM) tensor transport, implementing video decoding backends using OpenCV and FFmpeg in the multimodal crate, and updating the Qwen3-VL model spec to handle video pad replacements. Additionally, the Python gRPC servicer and Rust model gateway are updated to support SHM-based tensor transport and video preprocessor configurations. Feedback on the changes highlights a performance issue in the OpenCV video decoding loop, where failing to break early on a failed frame grab can lead to redundant and blocking operations.
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.
Collaborator
|
Hi @yechank-nvidia may you help fix the conflicts |
yechank-nvidia
force-pushed
the
yechan/mm-transport-opt
branch
from
June 21, 2026 15:09
25cc70d to
0e6246a
Compare
github-actions
Bot
added
tests
yechank-nvidia
force-pushed
the
yechan/mm-transport-opt
branch
from
June 21, 2026 15:22
0e6246a to
c8eb9ea
Compare
yechank-nvidia
marked this pull request as ready for review
June 21, 2026 15:22
yechank-nvidia
requested review from
CatherineSue,
key4ng and
slin1237
as code owners
June 21, 2026 15:22
Contributor
|
Hi @yechank-nvidia, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
needs-rebase
![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: c8eb9ea2d9
ℹ️ 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".
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 6
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
crates/multimodal/src/vision/processors/qwen_vl_base.rs (1)
746-747:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftUse PIL-exact bicubic in video resize paths when filter is CatmullRom.
Line 746 and Line 882 still use non-PIL resize paths for bicubic (
CatmullRom). That breaks the parity contract now enforced in image preprocessing and can shift video/video-rgb encoder inputs.Suggested direction
- let resized = resize(frame, tw32, th32, filter); + let resized = if filter == FilterType::CatmullRom { + resize_bicubic_pil(frame, tw32, th32) + } else { + resize(frame, tw32, th32, filter) + };- let resized = resize_rgb_bytes( - frame.data, frame.width, frame.height, tw32, th32, filter, - )?; + let resized = if filter == FilterType::CatmullRom { + // Convert bytes -> RGB image -> PIL-exact resize for parity path. + let src = image::RgbImage::from_raw(frame.width, frame.height, frame.data.to_vec()) + .ok_or_else(|| TransformError::ShapeError("invalid RGB frame buffer".to_string()))?; + resize_bicubic_pil(&DynamicImage::ImageRgb8(src), tw32, th32).to_rgb8() + } else { + resize_rgb_bytes(frame.data, frame.width, frame.height, tw32, th32, filter)? + };Also applies to: 882-889
🤖 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 `@crates/multimodal/src/vision/processors/qwen_vl_base.rs` around lines 746 - 747, The resize calls using the filter parameter at the locations where bicubic filtering (CatmullRom) is requested are not using PIL-exact bicubic resizing, which breaks parity with the image preprocessing contract. When the filter parameter is CatmullRom, replace the current resize function call with a PIL-exact bicubic resize implementation instead of using the generic resize path. Apply this fix to both occurrences of the resize calls (around line 746 in the main path and around line 882 in the alternate path) to ensure consistent bicubic filtering behavior across video and image preprocessing.
🤖 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 `@crates/multimodal/build.rs`:
- Line 5: The println statement in build.rs contains a hardcoded library path
that only works on x86_64 Debian/Ubuntu Linux systems. Replace this hardcoded
path with platform-aware logic that detects the target OS and architecture at
build time. Use cfg attributes (such as cfg(target_os = "...") and
cfg(target_arch = "...")) to conditionally set the correct library search path
for each supported platform, or alternatively use pkg-config to dynamically
locate the turbojpeg library. This ensures the build succeeds on macOS, Windows,
aarch64 Linux, and other supported platforms with appropriate paths for each
system.
In `@crates/multimodal/src/jpeg_turbo.rs`:
- Around line 62-63: The SAFETY comment in the unsafe block needs to be expanded
to explicitly document all the assumptions that make the unsafe operation sound.
Currently it covers the handle null check and buffer sizing, but you need to add
explicit explanations for: why the buf pointer is valid and safe to pass to the
FFI function (it comes from a valid, owned Vec), why there is no aliasing
concern during the FFI call, and how the return value check guarantees that no
out-of-bounds write occurred. Expand the existing SAFETY comment to
comprehensively cover all three of these invariants in addition to what is
already documented.
In `@crates/multimodal/tests/decode_preprocess_bench.rs`:
- Around line 18-20: The REAL_JPEG environment variable uses a hardcoded
machine-specific fallback path that causes unpredictable test failures outside
that particular environment. Replace the unwrap_or_else call on the REAL_JPEG
variable lookup with an expect call that requires the environment variable to be
explicitly set, following the same pattern used for PP_CONFIG, and provide a
clear error message instructing users to set the REAL_JPEG environment variable.
In `@model_gateway/src/routers/grpc/multimodal.rs`:
- Around line 1463-1471: The direct SHM writer path allocates a full Vec<u16>
before writing, causing excessive memory usage for large inputs. Instead of
collecting all conversions at once with `encoder_slice.iter().map(|&value|
convert(value)).collect()`, implement chunked processing similar to the
non-contiguous path by iterating through the encoder slice in manageable chunks,
converting each chunk to u16 values, and writing each chunk to the writer via
`writer.write_all(bytemuck::cast_slice(...))` separately. This approach keeps
peak memory bounded while maintaining the same functionality for the
little-endian target configuration.
In `@model_gateway/src/routers/grpc/proto_wrapper.rs`:
- Around line 431-438: The TokenSpeed SHM file creation in the OpenOptions
configuration uses create(true).truncate(true) which is vulnerable to symlink
attacks and file clobbering in world-writable directories. Replace
create(true).truncate(true) with create_new(true) to ensure exclusive file
creation that fails if the file already exists. Additionally, add explicit file
permission settings to restrict access to the owner only (using 0o600 on Unix
systems). This security fix should be applied to all SHM file creation
locations, including the probe file creation block starting around line 431 and
the other location around lines 507-510. Consider centralizing this SHM file
creation logic into a reusable function that consistently applies these security
measures across all instances.
In `@model_gateway/src/routers/grpc/utils/chat_utils.rs`:
- Around line 246-252: The pattern match in the String format branch is missing
audio_url handling. Currently, the match statement at line 246 only handles
image_url and video_url, causing audio_url to fall through to the _ case and be
silently dropped. Add audio_url to the pattern match alongside image_url and
video_url so that audio parts are also added to media_parts using the
image_placeholder, ensuring consistent handling with the OpenAI format branch at
line 271. This prevents silent data loss when messages contain audio content.
---
Outside diff comments:
In `@crates/multimodal/src/vision/processors/qwen_vl_base.rs`:
- Around line 746-747: The resize calls using the filter parameter at the
locations where bicubic filtering (CatmullRom) is requested are not using
PIL-exact bicubic resizing, which breaks parity with the image preprocessing
contract. When the filter parameter is CatmullRom, replace the current resize
function call with a PIL-exact bicubic resize implementation instead of using
the generic resize path. Apply this fix to both occurrences of the resize calls
(around line 746 in the main path and around line 882 in the alternate path) to
ensure consistent bicubic filtering behavior across video and image
preprocessing.
🪄 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: 45317c29-d6ab-477b-9b46-0e874c98b0ec
📒 Files selected for processing (16)
crates/grpc_client/proto/tokenspeed_scheduler.protocrates/multimodal/build.rscrates/multimodal/src/jpeg_turbo.rscrates/multimodal/src/lib.rscrates/multimodal/src/media.rscrates/multimodal/src/vision/processors/qwen_vl_base.rscrates/multimodal/src/vision/transforms.rscrates/multimodal/tests/decode_preprocess_bench.rscrates/multimodal/tests/preprocess_fingerprint.rscrates/multimodal/tests/resize_fingerprint.rsgrpc_servicer/smg_grpc_servicer/tokenspeed/servicer.pymodel_gateway/src/observability/metrics.rsmodel_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/multimodal.rsmodel_gateway/src/routers/grpc/proto_wrapper.rsmodel_gateway/src/routers/grpc/utils/chat_utils.rs
slin1237
requested changes
Jun 21, 2026
|
|
||
| // TokenSpeed multimodal tensor transport (shm vs inline) | ||
| describe_counter!( | ||
| "smg_tokenspeed_mm_tensors_total", |
Member
There was a problem hiding this comment.
we should not hardcode a metrics to a runtime specific. And runtime already has a label. Please change this
Collaborator
Author
There was a problem hiding this comment.
Thx for the comment.
Renamed the counters to neutral names (smg_mm_tensors_total, smg_mm_tensor_bytes_total, smg_mm_shm_write_failures_total) and moved the engine onto a runtime label, matching the existing smg_pd_* runtime-labeled metrics.
record_mm_tensor / record_mm_shm_write_failure now take a runtime arg; the TokenSpeed call sites pass "tokenspeed".
| _ => unreachable!("caller guarantees matching variant"), | ||
| }); | ||
| let req = client.build_generate_request_from_messages( | ||
| let shm_handles = tokenspeed_mm |
Member
There was a problem hiding this comment.
Same here. Only protocol can be engine specific. Everywhere else should be neutral. Tho there can be a case where one engine has a feature implemented and others don't. But the code should be neutral
Collaborator
Author
There was a problem hiding this comment.
Addressed both — metrics are now engine-neutral with a runtime label, and the TokenSpeed SHM lifecycle moved into the protocol layer so the dispatch is neutral.
Contributor
|
Hi @yechank-nvidia, this PR has merge conflicts that must be resolved before it can be merged. Please rebase your branch: git fetch origin main
git rebase origin/main
# resolve any conflicts, then:
git push --force-with-lease |
mergify
Bot
added
the
needs-rebase
Add counters so the shm-vs-inline transport split is observable:
- smg_tokenspeed_mm_tensors_total{path} tensors sent per path
- smg_tokenspeed_mm_tensor_bytes_total{path} bytes sent per path
- smg_tokenspeed_mm_shm_write_failures_total SHM writes that fell back to inline
Each tensor is metered exactly once: encoder inputs written directly to SHM are
counted at proto conversion (Shm storage arm); all other tensors (inline-storage
encoders + model_specific) are counted inside tokenspeed_tensor_payload. Exposed
on the existing Prometheus endpoint (:29000/metrics).
Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
Qwen2VL/Qwen3VL HF image processors default to BICUBIC (PIL resample=3) when the preprocessor config omits `resample`, but SMG's pil_to_filter fell back to bilinear (Triangle). Bilinear produces smoother encoder-input features; against vLLM on Qwen3.5-397B-A17B-NVFP4 / MMBench the SMG pixel_values matched HF only at corr 0.9994 (max abs diff 0.23). Pinning bicubic for the Qwen path raises the match to corr 0.99996 (max 0.04), aligning SMG with the reference HF/vLLM preprocessing. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
The eval harness sends multimodal content as [text, image]. SMG rendered the content parts in order, so the image landed AFTER the whole question — which measurably degrades VQA grounding (MMBench answers flip vs. image-first). vLLM always prepends media placeholders to the front (its default interleave_mm_strings=false); match that. transform_content_field now emits media (image/video/audio) before text in both content formats: OpenAI uses a stable partition; String collects placeholders first then text, joined by "\n". No-op for text-only/string content and content already media-first. A TODO(interleave) documents how to thread vLLM's interleave_mm_strings opt-out if ever needed. Verified: 10 unit tests plus an e2e render test against the real Qwen3.5 chat_template.jinja (image now at char 17 < question at char 60). End-to-end MMBench_DEV_EN_V11 (150-q sample) recovered 87% -> 96%. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
The Pillow-exact BICUBIC resize (resize_bicubic_pil), used for vLLM/PIL preprocessing parity, ran both passes single-threaded. On large images its scalar fixed-point arithmetic dominated preprocessing (≈4x slower than HF on 12MP). Each output row is an independent fixed-point integer sum, so band the horizontal/vertical passes across threads with std::thread::scope: identical arithmetic and inner-sum order => BIT-IDENTICAL output, no new dependency. Small images stay serial (par_threads gates on output size / rows-per-thread) to avoid spawn overhead. Accuracy is preserved unconditionally: resize_fingerprint.rs pins the exact byte output (fnv1a) of the serial implementation across up/downscale cases and asserts the parallel version reproduces it bit-for-bit. Golden parity (qwen35_parity) diffs vs HF are unchanged. Measured (Qwen3-VL preprocess, 224-core host): 12MP 862 -> 345 ms (2.5x), small/medium unchanged (resize not the hotspot there). Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
…rity Bring SMG's image path to bit-for-bit parity with the HF/PIL pipeline that vLLM uses, so vision-encoder inputs match and image accuracy is preserved: - jpeg_turbo: minimal libjpeg-turbo (TurboJPEG) FFI for RGB JPEG decode, matching Pillow's chroma-upsampling (the pure-Rust decoder differs by a few levels and shifts embeddings). build.rs links turbojpeg; media.rs routes JPEG decode through it (falling back to the `image` crate otherwise). - qwen_vl_base: Qwen3-VL preprocess (smart_resize, grid/token calc, fused normalize, patchify) producing HF-equivalent pixel_values + image_grid_thw. Pairs with the Pillow-exact BICUBIC resize already in transforms.rs. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
Profiling the preprocess on full-resolution images (this config does not downscale) showed the hotspot was normalize + patchify, not resize: 1920x1440 = normalize 23.9ms + patchify 26.2ms (resize skipped). Both stages are elementwise / per-block independent, so band them across threads with std::thread::scope -- identical arithmetic and write order => BIT-IDENTICAL: - deinterleave_rgb_to_planes (normalize): split the pixel range; each output element depends only on its own input byte. 7.5-20x faster. - patchify_into: split the (gt,pr,pc) blocks; each writes a contiguous, deterministic output region of pure copies (memory-bound, ~1.2-1.4x). par_threads gates on output size / rows-per-thread so small images stay serial. preprocess_fingerprint.rs pins the exact f32 encoder_input bytes and asserts the parallel path reproduces the serial output bit-for-bit (accuracy preserved unconditionally). decode_preprocess_bench.rs is the SMG-vs-HF harness. Measured (Qwen3-VL preprocess, 224-core host): 1280x960 15.5->9.0ms, 1920x1440 49.6->28.6ms, 12MP 862->173ms -- now faster than HF on 1280/12MP. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
…re REAL_JPEG - build.rs: derive the libturbojpeg search dir from the target arch (x86_64 / aarch64 Linux) and only add it when it exists, instead of hardcoding /usr/lib/x86_64-linux-gnu, which broke aarch64/macOS/Windows builds. The link-lib directive is always emitted; other platforms resolve via the default linker search path. (coderabbitai) - jpeg_turbo.rs: expand the SAFETY comment to cover handle lifetime, input buffer validity/no-aliasing, output buffer sizing, and the rc==0 guard. (coderabbitai) - decode_preprocess_bench.rs: require REAL_JPEG explicitly (like PP_CONFIG) instead of a machine-specific fallback path. (coderabbitai) Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
…ly abort - proto_wrapper.rs: create TokenSpeed SHM files with create_new(true) (no clobber/symlink in world-writable /dev/shm) and owner-only 0o600 mode (cfg(unix)), for both the writability probe and the payload writer. (coderabbitai, Major) - multimodal.rs / proto_wrapper.rs: assemble_tokenspeed now builds items imperatively and, if any item fails after its encoder input was serialized to SHM, unlinks the already-created /dev/shm segments (prior items + the pending tensor) before returning. Previously the partial TokenSpeedTensor::Shm handles were dropped without reaching the send-path cleanup, leaking files on repeated malformed/unsupported multimodal requests. (chatgpt-codex, P2) Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
…us path write_array_as_u16's contiguous fast path converted the whole bf16/f16 tensor into a single Vec<u16> before writing, so peak memory scaled with encoder-input size. Convert in CHUNK_VALUES-sized blocks (reusing one buffer) like the strided path already does, bounding peak conversion memory to ~512 KiB regardless of tensor size. (coderabbitai) Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
CI failed linking the Go/Python bindings: `cargo:rustc-link-lib=turbojpeg` made libturbojpeg a hard link-time dependency for every consumer of llm-multimodal, but the binding/CI runners don't ship it (`rust-lld: unable to find library -lturbojpeg`). The earlier build.rs path tweak didn't help — the link requirement itself was the problem. Load libturbojpeg via dlopen (libloading) at first use instead: try the runtime soname then the dev/macOS names, resolve the four TurboJPEG symbols, and cache the handle. No build script, no link-time dependency, so the crate and all consumers compile on any platform. Where the library is present (the serving image) decode still goes through it for PIL/vLLM parity; where it's absent decode_jpeg_rgb returns None and the caller falls back to the pure-Rust decoder. Removes build.rs and the #[link]/extern block; adds libloading. Verified the runtime path still decodes (3508.jpg 259x194, ~0.1ms/img). Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
…in1237)
Per review: SMG should stay engine-neutral; only the protocol layer may be
engine-specific.
- metrics: rename the runtime-baked counters smg_tokenspeed_mm_* ->
smg_mm_{tensors,tensor_bytes,shm_write_failures}_total and carry the engine as
a `runtime` label (matching the existing runtime-labeled metrics) instead of
the metric name. record_mm_tensor / record_mm_shm_write_failure take
`runtime`; TokenSpeed call sites pass "tokenspeed".
- client.rs: the per-engine build_chat/build_messages dispatch arms no longer
inline TokenSpeed SHM-handle collection + error-path cleanup. That
engine-specific SHM lifecycle moves into proto_wrapper::finish_tokenspeed_request
(the protocol layer), so the dispatch arms are thin like the other engines'.
No behavior change (88 grpc tests pass); metric series are renamed + gain a
runtime label.
Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
Clippy (deny warnings) flagged the new multimodal code: - use #[expect(clippy::too_many_arguments, reason=...)] instead of #[allow] for the band/patchify helpers (allow_attributes lint) - par_threads: .min(MAX).max(1) -> .clamp(1, MAX) (manual_clamp) - add `;` to the thread-spawn closures (semicolon_if_nothing_returned) - resize_bicubic_pil: gate the construction expect() with #[expect(clippy::expect_used, reason=...)] (buffer is sized by construction) - bring Metrics into scope at the two MM-metric call sites instead of fully qualifying (proto_wrapper, multimodal) - de-qualify HashMap and drop diagnostic eprintln! in the chat-template e2e test - fingerprint/bench tests: crate-level allow for unwrap/expect/print in tests; iterate to_le_bytes() by value No behavior change. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
- qwen_vl: route default-bicubic video frame resizing through the PIL-exact path (resize_bicubic_pil / new resize_bicubic_pil_rgb) in both preprocess_video and preprocess_video_rgb, matching the image path so video encoder inputs equal HF/vLLM bit-for-bit instead of diverging to the SIMD CatmullRom resizer. - grpc: do not infer a shared /dev/shm from TCP loopback. auto-mode SHM now requires a unix-domain-socket worker (proven same-host) or an explicit SMG_TOKENSPEED_SHM_ASSUME_LOOPBACK_SHARED operator assertion; loopback alone falls back to inline. Explicit shm mode is unchanged. - grpc: SHM cleanup only unlinks names carrying the smg-tokenspeed- prefix this transport creates, never arbitrary top-level /dev/shm entries. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
The String chat-content branch only matched image_url/video_url, so an audio_url part was silently dropped (no placeholder emitted) while the OpenAI branch already handled it. Match audio_url too, keeping the two content-format branches consistent and media-first. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
…MM transport env - MM tensor serialization (serialize_array, write_array_as_f32/u16, serialize_array_as_u16_bytes) only fast-paths C-contiguous arrays now. The as_slice_memory_order() fallback serialized Fortran-contiguous views in the wrong dimension order — the exact hazard serialize_array's own comment warned about. Non-C-contiguous arrays fall through to logical .iter(); inputs are C-contiguous in practice, so the wire bytes are unchanged. - Refresh the stale 'worker is local' comment in assemble_tokenspeed to match worker_shares_dev_shm (unix socket / operator-asserted loopback). - Document the TokenSpeed MM tensor transport env vars (transport mode, shm min bytes, loopback assertion, timing) plus the worker-side companions in the configuration reference. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
SMG_TOKENSPEED_TENSOR_TRANSPORT and SMG_TOKENSPEED_SHM_MIN_BYTES were introduced earlier in this same (unmerged) branch and then renamed to the SMG_TOKENSPEED_MM_* form for clarity, with the old names kept as 'legacy' fallbacks. Since neither name ever shipped on main, there is nothing to be backward-compatible with — drop the aliases and keep only the canonical SMG_TOKENSPEED_MM_* vars. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
auto-mode SHM previously inferred a shared /dev/shm from the worker URL and fell back to an operator env (SMG_TOKENSPEED_SHM_ASSUME_LOOPBACK_SHARED). URL locality only proves network reachability — a loopback/sidecar worker can have a private /dev/shm — so the env pushed a guess onto operators. Instead negotiate it (NIXL-style): the TokenSpeed worker advertises its /dev/shm filesystem identity (<boot_id>:<st_dev of /dev/shm>) in GetServerInfo's scheduler_info; discovery surfaces it as the worker's shm_namespace_id label; the router compares it to its own and enables SHM only on a match. boot_id pins the host; st_dev is the tmpfs superblock device, identical iff the same tmpfs backs both /dev/shm mounts — so it correctly detects sharing even across separate containers that share /dev/shm via --ipc/bind-mount (where mount namespaces differ but the superblock is the same). Any mismatch/missing token => inline. Verified on the live split-container deployment: a canary written to the router's /dev/shm is visible in the worker container, and both sides compute the identical token (<boot_id>:28) — i.e. token-match <=> actually-shared. (An earlier mount-namespace-inode token was rejected: it differs across those containers despite the shared tmpfs, which would wrongly force inline.) No new env, no proto change (reuses the existing scheduler_info Struct), and strictly more correct than URL inference. - client.rs: extract shm_namespace_id from scheduler_info into worker labels. - multimodal.rs: worker_shares_dev_shm compares the worker's token to the router's local <boot_id>:<st_dev>; remove url_is_loopback + the env; add a regression test that the local token resolves on Linux. - servicer.py: advertise _shm_namespace_id() in GetServerInfo. - docs: drop the loopback-assertion env; document the verified auto behavior. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
The video resize path added for HF/vLLM parity had no test: the existing video tests use 4x4 frames that never resize, so the new resize_bicubic_pil_rgb / CatmullRom branch was never exercised — while images already have fingerprint + golden bit-identity guards. - transforms: resize_bicubic_pil_rgb must equal the DynamicImage resize_bicubic_pil byte-for-byte (the core invariant making video frames match images/HF/vLLM), plus a wrong-length-buffer rejection test. - qwen_vl_base: with a frame that actually needs resizing (odd 7x9 -> factor- aligned target), preprocess_video and preprocess_video_rgb must produce bit-identical encoder inputs; asserts a resize was forced so the branch runs. Signed-off-by: yechank-nvidia <161688079+yechank-nvidia@users.noreply.github.com>
slin1237
force-pushed
the
yechan/mm-transport-opt
branch
from
June 22, 2026 12:47
95e0dd6 to
1c766de
Compare
mergify
Bot
removed
the
needs-rebase
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1c766de148
ℹ️ 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".
| #[cfg(unix)] | ||
| { | ||
| use std::os::unix::fs::OpenOptionsExt; | ||
| opts.mode(0o600); |
There was a problem hiding this comment.
In deployments where the router and TokenSpeed worker share /dev/shm but run as different non-root UIDs, auto mode can still choose SHM because worker_shares_dev_shm only compares the tmpfs identity. These payload files are created owner-only, and the worker later opens the path from the servicer/scheduler side, so a different UID gets EACCES and the request fails instead of falling back to inline; either include readable ownership/UID in the handshake or create the segment with permissions the verified worker can read.
Useful? React with 👍 / 👎.
2 tasks
slin1237
added a commit
that referenced
this pull request
Jun 23, 2026
`test_multi_images_mixed` sends the same pug twice (mixed base64 + URL) and asserted the model counts "3". Engines legitimately differ on byte-identical multimodal inputs: vLLM deduplicates them (encodes the duplicate once), sglang keeps both — so the literal count is engine-dependent. This surfaced after #1604 made image decode bit-deterministic, so identical inputs now produce identical pixel tensors that vLLM dedups (sglang is unaffected). Drop the exact-count assertion and keep the duplicate-detection assertion (which both engines satisfy) plus the dog/pug content checks, so the test still validates multi-image + mixed base64/URL + duplicate handling. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
This was referenced Jul 8, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Problem
Two multimodal issues on the SMG ↔ TokenSpeed path:
Tensor-transport overhead. Large precomputed multimodal encoder inputs (image/video
pixel_values/embeddings) are serialized inline over gRPC, adding serialization + memory overhead under high-concurrency image/video load, and shared-memory payloads had no cleanup if a request failed mid-construction.Image preprocessing diverged from the HF/PIL pipeline vLLM uses, which cost accuracy and speed:
[text, image]; SMG rendered the image after the whole question (content order), TokenSpeed was ~7pp below vLLM on MMBench(image).Solution
std::thread::scope— bit-identical (each output element/block is independent; serial fallback for small images).Changes
Transport
grpc_servicer/.../tokenspeed/servicer.py: SHM tensor read +ShmTensorHandle, unlink-after-read, validation, MM timing logsmodel_gateway/.../grpc/{multimodal,proto_wrapper,client}.rs: SHM-capable tensor storage, threshold/env controls, reduced-dtype payloads, failure cleanupmodel_gateway/.../observability/metrics.rs:smg_*MM-transport metricscrates/grpc_client/proto/tokenspeed_scheduler.proto: payload field(s)Preprocessing parity & accuracy
model_gateway/.../grpc/utils/chat_utils.rs(transform_content_field): media-first content ordering for String + OpenAI formats;TODO(interleave)forinterleave_mm_stringsopt-outcrates/multimodal/src/jpeg_turbo.rs(+build.rs,lib.rs,media.rs): libjpeg-turbo RGB decode FFI (PIL-compatible)crates/multimodal/src/vision/transforms.rs: Pillow-exactresize_bicubic_pil+ parallelized resize/normalizecrates/multimodal/src/vision/processors/qwen_vl_base.rs: HF-equivalent Qwen3-VL preprocess + parallelized patchifyTests
tests/resize_fingerprint.rs,tests/preprocess_fingerprint.rs: bit-identity guards (fnv) — parallel output == serial, byte-for-bytetests/decode_preprocess_bench.rs: SMG-vs-HF decode/preprocess microbenchTest Plan
Accuracy (MMBench_DEV_EN_V11, 397B-A17B-NVFP4, 8×B200): image-first fix recovers TS 87% → 96% (150-q sample, 0 errors), matching/exceeding vLLM(94%); e.g. idx 3508 (key C) flips from wrong
A(image-last) to correctC.Parity:
resize_bicubic_pil_bit_identity+preprocess_bit_identityassert the parallel path reproduces the serial output bit-for-bit; golden parity vs HF unchanged.Preprocess speed (Qwen3-VL, 224-core host): decode at parity with PIL (both libjpeg-turbo); preprocess 1280×960 15.5→9.0 ms, 1920×1440 49.6→28.6 ms, 12MP 862→173 ms — faster than HF on 1280²/12MP, all bit-identical.
Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
Release Notes
resize_to_fit.