Repository navigation
feat(multimodal): optimize EPD encode routing - #1853
Conversation
📝 WalkthroughWalkthroughAdds ChangesRDMA Multimodal Transport
Estimated code review effort: 4 (Complex) | ~75 minutes Vision Preprocessing Buffer Pooling
Estimated code review effort: 3 (Moderate) | ~30 minutes Host-DRAM Pixel Preprocessing Cache
Estimated code review effort: 3 (Moderate) | ~30 minutes Possibly related issues
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ 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
dependencies
![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 introduces support for Encode-Prefill-Decode (EPD) disaggregated mode, adding a new TokenSpeedEncoder gRPC service, EPD routing in the model gateway, an optional NIXL-based RDMA pixel transport, a host-DRAM LRU cache for preprocessed encoder inputs, and recycling pools for large vision buffers. The review feedback highlights several critical issues: an unawaited coroutine when sending payloads to the scheduler, a need to use tokio::task::JoinSet instead of futures::future::join_all for concurrent tasks, the importance of offloading expensive DynamicImage cloning to spawn_blocking to avoid blocking async executor threads, and a thread-safety issue requiring a lock for the shared _LOCAL_IP_CACHE dictionary.
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.
Optimize multimodal tensor handling for EPD encode routing. Add optional fast-path transfer support, scratch buffer reuse, and model preprocessing/serialization improvements. Signed-off-by: chenht2022 <chenht2022@gmail.com>
…rtMode Fold the multimodal pixel RDMA lane into the engine-agnostic transport config layer instead of a standalone env gate. Add `TransportMode::Rdma` so it is selectable via `--multimodal-tensor-transport rdma` / `SMG_MM_TENSOR_TRANSPORT=rdma` alongside inline/shm/auto; the RDMA gate now consults the resolved transport mode, keeping the legacy `SMG_MM_PIXEL_RDMA` env as a backward-compatible fallback. - openai-protocol: TransportMode::Rdma variant + parse/as_str/Display/FromStr - transport: resolve_mm_shm_enabled treats rdma as non-SHM; feature-gated mm_default_transport_is_rdma() accessor the gate consults - mm_rdma: rdma_enabled() honors TransportMode::Rdma (mm-rdma feature) - CLI value_parser, Python arg, binding, and docs recognize the rdma value Co-authored-by: chenht2022 <chenht2022@gmail.com> Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
github-actions
Bot
added
documentation
Post-rebase lint fixes so smg-project#1853 passes CI: - import crate::vision::scratch (kimi_k25, transforms) instead of calling it through an absolute path (clippy::absolute_paths, denied in this repo) - shorten the rdma transport-mode help string under the 120-col ruff limit Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
![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: f0eb4f010d
ℹ️ 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".
| if os.environ.get("SMG_MM_PIXEL_RDMA") not in ("1", "true"): | ||
| return |
There was a problem hiding this comment.
When the gateway is enabled via the new first-class --multimodal-tensor-transport rdma / SMG_MM_TENSOR_TRANSPORT=rdma path, it can export TokenSpeed encoder inputs as remote TensorData, but this worker-side puller only initializes when the legacy SMG_MM_PIXEL_RDMA variable is set. In that configuration the first remote payload hits feature_from_remote() with no landing pool and the request fails instead of pulling or falling back inline, so the new transport option is unusable unless operators also know to set the old env var on every TokenSpeed worker.
Useful? React with 👍 / 👎.
Make smg-project#1853 pass CI (clippy --all-targets --all-features -D warnings, ruff): - absolute_paths: import crate::vision::scratch (kimi_k25, transforms) and crate::routers::grpc::mm_rdma (proto_wrapper) instead of absolute-path calls - allow_attributes: #[expect(unsafe_code, reason=...)] not #[allow] for the NIXL FFI unsafe blocks (nixl.rs) - ruff: shorten the rdma transport-mode help string under 120 cols Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
There was a problem hiding this comment.
Pull request overview
This PR optimizes the multimodal TokenSpeed Encode/Prefill/Decode (EPD) routing path by introducing an opt-in RDMA/NIXL “remote” tensor payload for large pixel buffers, plus several CPU-side preprocessing and serialization optimizations to reduce allocations and copy overhead.
Changes:
- Add an optional NIXL RDMA “remote” payload for TokenSpeed encoder inputs (gateway export + worker pull), including new
TransportMode::Rdma. - Add a host-DRAM LRU pixel cache for single-image preprocessing and refactor preprocessing to reuse cached results when enabled.
- Reduce vision preprocessing allocation churn via scratch-buffer reuse and hot-path improvements in model-specific processors.
Reviewed changes
Copilot reviewed 29 out of 29 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| model_gateway/src/routers/grpc/proto_wrapper.rs | Adds TokenSpeed tensor “remote” storage and exports inline encoder inputs to NIXL descriptors. |
| model_gateway/src/routers/grpc/multimodal/transport.rs | Extends transport resolution to support rdma and exposes a mm_default_transport_is_rdma helper (feature-gated). |
| model_gateway/src/routers/grpc/multimodal/serialize.rs | Updates serialization tests to account for the new Remote tensor storage variant. |
| model_gateway/src/routers/grpc/multimodal/process.rs | Adds single-image pixel-cache path and factors out image preprocessing helpers. |
| model_gateway/src/routers/grpc/multimodal/pixel_cache.rs | Introduces a byte-budgeted LRU cache for per-image preprocessed encoder inputs. |
| model_gateway/src/routers/grpc/multimodal/mod.rs | Wires pixel_cache module and exports mm_default_transport_is_rdma under mm-rdma. |
| model_gateway/src/routers/grpc/multimodal/config.rs | Plumbs an optional PixelCache into multimodal components (env-configured). |
| model_gateway/src/routers/grpc/mod.rs | Adds the mm_rdma module to the gRPC router. |
| model_gateway/src/routers/grpc/mm_rdma/stub.rs | Provides a no-op RDMA shim when mm-rdma feature is disabled. |
| model_gateway/src/routers/grpc/mm_rdma/nixl.rs | Implements gateway-side NIXL export with a pre-registered arena + slot lease/reap logic. |
| model_gateway/src/routers/grpc/mm_rdma/mod.rs | Feature-gates the real NIXL implementation vs the stub. |
| model_gateway/src/routers/grpc/epd_encode.rs | Exports EPD encode pixel payloads to NIXL using the encode bootstrap room. |
| model_gateway/src/routers/grpc/client.rs | Exports non-EPD TokenSpeed encoder inputs to NIXL remote descriptors before proto conversion. |
| model_gateway/src/observability/metrics.rs | Extends MM tensor metrics labeling to include “remote”. |
| model_gateway/src/main.rs | Updates CLI validation messaging to include rdma. |
| model_gateway/src/config/types.rs | Updates router config docs to include rdma transport. |
| model_gateway/Cargo.toml | Adds mm-rdma feature + optional nixl-sys and adds lru dependency. |
| grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py | Adds receiver-side remote payload support via a shared RDMA pixel puller. |
| grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py | Adds a persistent NIXL READ agent + landing pool to pull remote pixel tensors. |
| grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py | Adds encode-worker-side remote payload handling (room-validated) with optional SHM publish. |
| docs/reference/configuration.md | Documents rdma as a multimodal tensor transport mode and its requirements/fallback. |
| crates/protocols/src/worker.rs | Extends TransportMode enum parsing/display to include Rdma. |
| crates/multimodal/src/vision/transforms.rs | Switches a hot per-image CHW buffer allocation to use the new scratch pool. |
| crates/multimodal/src/vision/scratch.rs | Introduces thread-local scratch-buffer pools for large vision allocations. |
| crates/multimodal/src/vision/processors/qwen3_vl.rs | Adds a per-image-independence regression test supporting the gateway pixel-cache assumption. |
| crates/multimodal/src/vision/processors/kimi_k25.rs | Reduces allocations by reusing scratch buffers and appending patches directly into a pooled batch buffer. |
| crates/multimodal/src/vision/mod.rs | Exposes the new scratch module internally. |
| bindings/python/src/smg/router_args.py | Extends Python CLI arg choices/docs to include rdma. |
| bindings/python/src/lib.rs | Updates Python binding error messages to include rdma. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| cache.insert( | ||
| key, | ||
| Arc::new(CachedPreprocessedItem { | ||
| preprocessed: preprocessed.clone(), | ||
| }), | ||
| ); | ||
| Ok(preprocessed) |
| // Patchify directly into the pooled batch buffer, then recycle the | ||
| // CHW tensor's storage (standard layout, offset 0) for the next image. | ||
| Self::extract_patches_into(&tensor, self.patch_size, &mut all_patches); | ||
| let (storage, _offset) = tensor.into_raw_vec_and_offset(); | ||
| scratch::give_f32(storage); |
| agent = self._nixl_agent | ||
| if agent is None or self._landing_free is None: | ||
| raise RuntimeError(f"{self._log_prefix}: remote payload but landing pool unavailable") |
| fn build_arena(g: &Mutex<GatewayRdma>) -> Result<SlotArena, nixl_sys::NixlError> { | ||
| let n_slots = pool_slots(); | ||
| let slot_bytes = slot_bytes(); | ||
| let total = n_slots.saturating_mul(slot_bytes); | ||
| // Leak the arena for the process lifetime: `base` must stay registered + stable, | ||
| // and the gateway agent never shuts down. Holding it as raw memory (not a typed | ||
| // Box) is what makes the per-slot memcpy through `base` sound. | ||
| let boxed = vec![0u8; total].into_boxed_slice(); | ||
| let base = Box::into_raw(boxed) as *mut u8 as usize; |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4921987825
ℹ️ 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".
| pub(crate) fn mm_default_transport_is_rdma() -> bool { | ||
| mm_transport_defaults().mode == TransportMode::Rdma |
There was a problem hiding this comment.
This predicate is what mm_rdma::rdma_enabled() uses before exporting pixels, but it only checks the router-level default and never the WorkerSpec.multimodal_tensor_transport override that the rest of transport resolution gives priority to. In a deployment with global inline and an encode/single worker set to rdma, the new path silently stays inline; conversely a global rdma setting still exports remote payloads to a worker explicitly overridden to inline/shm. Pass the selected worker/mode through the RDMA export path instead of gating on the process default only.
Useful? React with 👍 / 👎.
| pp_config: PreProcessorConfig, | ||
| images: &[Arc<ImageFrame>], | ||
| ) -> Result<PreprocessedEncoderInputs> { | ||
| let raw_images: Vec<image::DynamicImage> = images.iter().map(|f| f.image.clone()).collect(); |
There was a problem hiding this comment.
When SMG_MM_PIXEL_CACHE_MB is enabled and a single-image request misses the cache, this clones the decoded image bytes before entering spawn_blocking, so large images are copied on a Tokio runtime thread. The non-cache path deliberately performs the same clone inside the blocking closure to avoid stalling unrelated requests; this new cache-miss path should move the ImageFrame/image clone into the blocking task as well.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 8
🤖 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/src/vision/scratch.rs`:
- Around line 1-58: Add direct unit tests for the scratch buffer pool contract
in scratch.rs: verify take_f32 returns a zero-filled Vec even after reusing a
previously populated buffer, and verify give_f32 followed by take_f32_cap
round-trips the same-thread pooled storage with capacity preserved. Use the
existing take_f32, take_f32_cap, and give_f32 helpers to keep the test focused
on the pooling behavior and the MAX_THREAD_POOLED cap.
In `@grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py`:
- Around line 135-158: The landing-pool initialization in rdma_pixel.py
allocates a very large pinned buffer by default, so verify that the current
SMG_RDMA_LANDING_SLOTS and SMG_RDMA_SLOT_BYTES values are intentionally large
enough for all servicers. If not, reduce the defaults in the landing-pool setup
and update any related docs or config guidance so users know the expected memory
footprint when SMG_MM_PIXEL_RDMA is enabled. Make the change around the
self._landing allocation and the slot_bytes/n_slots defaults.
In `@grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py`:
- Around line 1038-1046: The remote multimodal reconstruction path in
_mm_inputs_from_itemized_proto is still running on the asyncio loop and can
block RPC handling when RdmaPixelPuller.feature_from_remote is waiting on slots,
metadata, or READ polling. Offload the remote branch the same way the encoder
path does by using asyncio.to_thread around the feature_from_remote call, and
keep Generate/_build_generate_req from doing this work inline on the event loop.
Make sure any shared puller state or request mutation used by
RdmaPixelPuller.feature_from_remote remains thread-safe when moved off-thread.
In `@model_gateway/src/routers/grpc/mm_rdma/nixl.rs`:
- Around line 289-308: The RDMA fallback in agent() and build_arena() assumes
nixl_sys will return Err values, but the stub can abort on first call instead of
erroring. Verify the pinned nixl-sys behavior for init_agent(), build_arena(),
and any first-touch NIXL calls, and if abort is possible add a startup
preflight/probe before enabling RDMA so the code disables mm-rdma gracefully
instead of reaching a crash-prone path. Use the existing agent() and
build_arena() flow to gate initialization based on that check.
- Around line 423-454: The reaper thread spawn failure is being swallowed in
spawn_reaper by discarding the Result from std::thread::Builder::spawn, so
change it to handle the error path explicitly and log the failure with enough
context to diagnose why the epd-rdma-reaper thread was not created. Keep the
existing thread setup and loop intact, but replace the silent .ok() handling
with error reporting tied to spawn_reaper/Builder::spawn so OS resource
exhaustion or similar startup failures are visible in logs.
In `@model_gateway/src/routers/grpc/multimodal/process.rs`:
- Around line 440-458: The cache path in process_image_batch is still cloning
the full PreprocessedEncoderInputs on both cache hits and inserts. Update
CachedPreprocessedItem and the downstream
PrecomputedMultimodalIntermediate.preprocessed ownership model so they store
Arc<PreprocessedEncoderInputs> instead of an owned value, then have
process_image_batch return/clones Arc handles rather than deep-copying
encoder_input on every hit or cache fill.
- Around line 426-459: The cache-aside flow in preprocess_image_cached still
allows concurrent requests for the same PixelCacheKey to run
preprocess_image_batch multiple times before insert completes. Add in-flight
request coalescing around the existing cache lookup/compute path so only one
task computes per key and the others await the same result, using a
single-flight mechanism keyed by PixelCacheKey (for example with DashMap plus
OnceCell or equivalent). Keep the current cache.get and cache.insert behavior,
but wrap the miss path so duplicate misses for the same image_hash and
config_fingerprint share one preprocessing future.
- Around line 461-479: `preprocess_image_batch` is cloning each `DynamicImage`
before `spawn_blocking`, so the expensive pixel-buffer copy still runs on the
async runtime thread. Update this function to follow the same pattern used in
`process_multimodal_parts` by only collecting the cheap `Arc<ImageFrame>`
references upfront and performing `f.image.clone()` inside the blocking closure.
Keep the `VisionProcessorRegistry::find` and `processor.preprocess` calls in the
blocking task so the full preprocessing work stays off the tokio worker thread.
🪄 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: df90e7fd-5ce0-4ef9-9a47-679549bccd87
📒 Files selected for processing (29)
bindings/python/src/lib.rsbindings/python/src/smg/router_args.pycrates/multimodal/src/vision/mod.rscrates/multimodal/src/vision/processors/kimi_k25.rscrates/multimodal/src/vision/processors/qwen3_vl.rscrates/multimodal/src/vision/scratch.rscrates/multimodal/src/vision/transforms.rscrates/protocols/src/worker.rsdocs/reference/configuration.mdgrpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.pygrpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.pygrpc_servicer/smg_grpc_servicer/tokenspeed/servicer.pymodel_gateway/Cargo.tomlmodel_gateway/src/config/types.rsmodel_gateway/src/main.rsmodel_gateway/src/observability/metrics.rsmodel_gateway/src/routers/grpc/client.rsmodel_gateway/src/routers/grpc/epd_encode.rsmodel_gateway/src/routers/grpc/mm_rdma/mod.rsmodel_gateway/src/routers/grpc/mm_rdma/nixl.rsmodel_gateway/src/routers/grpc/mm_rdma/stub.rsmodel_gateway/src/routers/grpc/mod.rsmodel_gateway/src/routers/grpc/multimodal/config.rsmodel_gateway/src/routers/grpc/multimodal/mod.rsmodel_gateway/src/routers/grpc/multimodal/pixel_cache.rsmodel_gateway/src/routers/grpc/multimodal/process.rsmodel_gateway/src/routers/grpc/multimodal/serialize.rsmodel_gateway/src/routers/grpc/multimodal/transport.rsmodel_gateway/src/routers/grpc/proto_wrapper.rs
| //! Recycling pools for the large per-image vision buffers. | ||
| //! | ||
| //! The preprocess pipeline allocates tens of MB per image (the [C, H, W] f32 | ||
| //! tensor and the batched patch buffer are each large). | ||
| //! Freshly-allocated buffers of this size bypass the allocator's reuse paths | ||
| //! (glibc caps non-main-arena chunks at 64 MB and mmaps anything larger or | ||
| //! colder), so every image pays tens of thousands of minor page faults; the | ||
| //! fault path serializes process-wide and caps the data plane's effective | ||
| //! parallelism. Recycling keeps the pages mapped and hot. | ||
| //! | ||
| //! A lock-free thread-local pool serves same-thread take/give (preprocess | ||
| //! internals run on blocking-pool threads). The pool is capped to bound | ||
| //! residency; buffers beyond the cap are dropped. | ||
|
|
||
| use std::cell::RefCell; | ||
|
|
||
| /// Max recycled buffers kept per thread per class; excess is dropped. The | ||
| /// vision path holds at most a couple of live tensors per request, so a small | ||
| /// cap captures same-thread reuse. | ||
| const MAX_THREAD_POOLED: usize = 2; | ||
|
|
||
| thread_local! { | ||
| static F32_LOCAL: RefCell<Vec<Vec<f32>>> = const { RefCell::new(Vec::new()) }; | ||
| } | ||
|
|
||
| macro_rules! pool_impl { | ||
| ($take_cap:ident, $give:ident, $ty:ty, $local:ident) => { | ||
| /// Take an empty `Vec` with at least `cap` capacity, reusing pooled storage. | ||
| pub fn $take_cap(cap: usize) -> Vec<$ty> { | ||
| let mut v = $local.with(|p| p.borrow_mut().pop()).unwrap_or_default(); | ||
| v.clear(); | ||
| v.reserve(cap); | ||
| v | ||
| } | ||
|
|
||
| /// Return a buffer for reuse by a later same-thread take. | ||
| pub fn $give(v: Vec<$ty>) { | ||
| if v.capacity() == 0 { | ||
| return; | ||
| } | ||
| $local.with(|p| { | ||
| let mut p = p.borrow_mut(); | ||
| if p.len() < MAX_THREAD_POOLED { | ||
| p.push(v); | ||
| } | ||
| }); | ||
| } | ||
| }; | ||
| } | ||
|
|
||
| pool_impl!(take_f32_cap, give_f32, f32, F32_LOCAL); | ||
|
|
||
| /// Take a zero-filled `Vec<f32>` of exactly `len`, reusing pooled storage. | ||
| pub fn take_f32(len: usize) -> Vec<f32> { | ||
| let mut v = take_f32_cap(len); | ||
| v.resize(len, 0.0); | ||
| v | ||
| } |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Solid pooling implementation; consider adding direct unit tests for the take/give contract.
The take/clear/reserve/resize sequence is correct: take_f32_cap always returns a Vec with capacity ≥ cap, and take_f32's resize(len, 0.0) guarantees the returned buffer is fully zero-initialized regardless of stale pooled contents, so there's no uninitialized-data leak across reuse cycles. give_f32 correctly rejects zero-capacity buffers and respects the MAX_THREAD_POOLED cap.
This module currently has no dedicated unit tests (correctness is only exercised indirectly via kimi_k25.rs/transforms.rs consumers and the qwen3_vl.rs equivalence test). A small test verifying take_f32 zero-fills and that give_f32/take_f32_cap round-trip correctly would guard this shared primitive against regressions independent of any specific processor.
🤖 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/scratch.rs` around lines 1 - 58, Add direct unit
tests for the scratch buffer pool contract in scratch.rs: verify take_f32
returns a zero-filled Vec even after reusing a previously populated buffer, and
verify give_f32 followed by take_f32_cap round-trips the same-thread pooled
storage with capacity preserved. Use the existing take_f32, take_f32_cap, and
give_f32 helpers to keep the test focused on the pooling behavior and the
MAX_THREAD_POOLED cap.
| slot_bytes = int(os.environ.get("SMG_RDMA_SLOT_BYTES", 32 * 1024 * 1024)) | ||
| n_slots = int(os.environ.get("SMG_RDMA_LANDING_SLOTS", 64)) | ||
| self._landing = torch.empty( | ||
| n_slots * slot_bytes, | ||
| dtype=torch.uint8, | ||
| pin_memory=True, | ||
| ) | ||
| self._landing_np = self._landing.numpy() | ||
| self._landing_base = self._landing.data_ptr() | ||
| self._landing_slot_bytes = slot_bytes | ||
| reg = self._nixl_agent.get_reg_descs( | ||
| [(self._landing_base, n_slots * slot_bytes, 0, "")], | ||
| "DRAM", | ||
| ) | ||
| self._nixl_agent.register_memory(reg) | ||
| self._landing_free = queue.Queue() | ||
| for i in range(n_slots): | ||
| self._landing_free.put(i) | ||
| logger.info( | ||
| "%s: puller up (landing %d slots x %d B)", | ||
| self._log_prefix, | ||
| n_slots, | ||
| slot_bytes, | ||
| ) |
There was a problem hiding this comment.
🧹 Nitpick | 🔵 Trivial
Verify the default landing-pool footprint is acceptable.
With defaults SMG_RDMA_LANDING_SLOTS=64 and SMG_RDMA_SLOT_BYTES=32MiB, torch.empty(n_slots * slot_bytes, pin_memory=True) allocates ~2 GiB of pinned host memory per servicer at init whenever SMG_MM_PIXEL_RDMA is enabled. Pinned pages are non-swappable and reduce page-cache/headroom; on nodes that run both the encoder and scheduler servicers this doubles. Please confirm the default is intended, or consider a smaller default with docs guidance.
🤖 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/tokenspeed/rdma_pixel.py` around lines 135 -
158, The landing-pool initialization in rdma_pixel.py allocates a very large
pinned buffer by default, so verify that the current SMG_RDMA_LANDING_SLOTS and
SMG_RDMA_SLOT_BYTES values are intentionally large enough for all servicers. If
not, reduce the defaults in the landing-pool setup and update any related docs
or config guidance so users know the expected memory footprint when
SMG_MM_PIXEL_RDMA is enabled. Make the change around the self._landing
allocation and the slot_bytes/n_slots defaults.
| if encoder_input.WhichOneof("payload") == "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) |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Confirm _build_generate_req is invoked synchronously on the loop (no to_thread offload).
rg -nP -C3 '_build_generate_req|to_thread' grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.pyRepository: lightseekorg/smg
Length of output: 800
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the Generate path and the remote feature pull implementation.
sed -n '130,220p' grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py
printf '\n====\n'
sed -n '1030,1075p' grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py
printf '\n====\n'
rg -n -C3 'feature_from_remote|to_thread|asyncio\.to_thread|run_in_executor|time\.sleep|SMG_RDMA_' grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.pyRepository: lightseekorg/smg
Length of output: 7205
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the relevant symbols and inspect the puller implementation.
ast-grep outline grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py --view expanded >/tmp/servicer_outline.txt
cat /tmp/servicer_outline.txt | sed -n '1,220p'
printf '\n====\n'
rg -n -C4 'class RdmaPixelPuller|def feature_from_remote|feature_from_remote\(' -S .
printf '\n====\n'
# Read the implementation file(s) if found by the search above.Repository: lightseekorg/smg
Length of output: 6173
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the RDMA puller implementation and the encoder-servicer offload pattern.
sed -n '1,260p' grpc_servicer/smg_grpc_servicer/tokenspeed/rdma_pixel.py
printf '\n====\n'
sed -n '90,160p' grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.py
printf '\n====\n'
rg -n -C4 'asyncio\.to_thread|run_in_executor|_parse_and_pickle|feature_from_remote|thread|lock|queue' grpc_servicer/smg_grpc_servicer/tokenspeed/{encoder_servicer.py,rdma_pixel.py,servicer.py}Repository: lightseekorg/smg
Length of output: 26673
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect request-building and multimodal reconstruction to see what state is touched.
sed -n '811,1018p' grpc_servicer/smg_grpc_servicer/tokenspeed/servicer.py
printf '\n====\n'
sed -n '1,120p' grpc_servicer/smg_grpc_servicer/tokenspeed/encoder_servicer.pyRepository: lightseekorg/smg
Length of output: 14435
Move remote multimodal reconstruction off the asyncio loop Generate() calls _build_generate_req() inline, and the remote MM path inside _mm_inputs_from_itemized_proto() reaches RdmaPixelPuller.feature_from_remote(), which can block on landing-slot acquisition, metadata handshake, and the READ poll loop. That can stall other RPCs on this servicer; mirror the encoder path’s asyncio.to_thread(...) offload and keep the shared puller/request mutation thread-safe.
🤖 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/tokenspeed/servicer.py` around lines 1038 -
1046, The remote multimodal reconstruction path in
_mm_inputs_from_itemized_proto is still running on the asyncio loop and can
block RPC handling when RdmaPixelPuller.feature_from_remote is waiting on slots,
metadata, or READ polling. Offload the remote branch the same way the encoder
path does by using asyncio.to_thread around the feature_from_remote call, and
keep Generate/_build_generate_req from doing this work inline on the event loop.
Make sure any shared puller state or request mutation used by
RdmaPixelPuller.feature_from_remote remains thread-safe when moved off-thread.
| fn agent() -> Option<&'static Mutex<GatewayRdma>> { | ||
| AGENT | ||
| .get_or_init(|| { | ||
| if !rdma_enabled() { | ||
| return None; | ||
| } | ||
| match init_agent() { | ||
| Ok(g) => { | ||
| spawn_reaper(); | ||
| debug!("EPD RDMA: gateway NIXL agent + UCX backend up"); | ||
| Some(Mutex::new(g)) | ||
| } | ||
| Err(e) => { | ||
| error!(error = ?e, "EPD RDMA: agent init failed; falling back to inline pixels"); | ||
| None | ||
| } | ||
| } | ||
| }) | ||
| .as_ref() | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🔵 Trivial | 🏗️ Heavy lift
NIXL init/registration is Result-based, but stub-api may abort instead of erroring.
agent() and build_arena() treat nixl_sys failures as recoverable (Err(e) => None, logged and falling back to inline). If the linked nixl-sys stub implementation aborts the process on the first real NIXL call rather than returning an Err (see the crate documentation, and this file's own comment at the top acknowledging "a missing libnixl_capi.so aborts on the first stub call, not a no-op"), then this fallback path can never actually engage — the whole gateway process would crash on first RDMA export attempt in an environment where mm-rdma is enabled + SMG_RDMA_LISTEN_IP is set but the native library isn't installed. This is already flagged as a known caveat in the module doc comment, so it appears to be an accepted tradeoff, but it's worth confirming the exact abort/dlopen behavior for the pinned nixl-sys = "0.10.1" version and considering a startup preflight (e.g., dlopen probe) that disables RDMA gracefully instead of enabling a path that can abort under load.
Also applies to: 334-356
🤖 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/mm_rdma/nixl.rs` around lines 289 - 308, The
RDMA fallback in agent() and build_arena() assumes nixl_sys will return Err
values, but the stub can abort on first call instead of erroring. Verify the
pinned nixl-sys behavior for init_agent(), build_arena(), and any first-touch
NIXL calls, and if abort is possible add a startup preflight/probe before
enabling RDMA so the code disables mm-rdma gracefully instead of reaching a
crash-prone path. Use the existing agent() and build_arena() flow to gate
initialization based on that check.
| fn spawn_reaper() { | ||
| std::thread::Builder::new() | ||
| .name("epd-rdma-reaper".into()) | ||
| .spawn(|| loop { | ||
| std::thread::sleep(REAPER_TICK); | ||
| let (Some(g), Some(a)) = ( | ||
| AGENT.get().and_then(|o| o.as_ref()), | ||
| ARENA.get().and_then(|o| o.as_ref()), | ||
| ) else { | ||
| continue; | ||
| }; | ||
| if let Ok(mut notifs) = NotificationMap::new() { | ||
| { | ||
| let guard = g.lock(); | ||
| let _ = guard.agent.get_notifications(&mut notifs, None); | ||
| } | ||
| if let Ok(map) = notifs.take_notifs() { | ||
| for (_agent, tags) in map { | ||
| for tag in tags { | ||
| if let Ok(room) = tag.parse::<i64>() { | ||
| a.pool.free_room(room); | ||
| } | ||
| } | ||
| } | ||
| } | ||
| } | ||
| // TTL sweep: reclaim slots whose READ-notif never arrived. `slot_ttl()` is | ||
| // derived to exceed the worker's max hold, so this never races a live READ. | ||
| let _ = a.pool.reap_stale(Instant::now(), slot_ttl()); | ||
| }) | ||
| .ok(); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Reaper thread spawn failure is silently swallowed.
spawn_reaper() discards the Result from std::thread::Builder::spawn(...) via .ok() with no logging. If the OS fails to spawn this thread (e.g., resource exhaustion), both the free-notif draining and the TTL sweep never run again for the life of the process — every leased slot leaks permanently, the pool silently exhausts, and every subsequent export_pixel_buffer call fails closed to inline with no diagnostic trail pointing at the actual root cause. This is a materially different failure mode than the "lost single notif" scenario the TTL sweep is designed to guard against, since here the whole recycling mechanism is dead from the start.
🩹 Proposed fix: log the spawn failure
fn spawn_reaper() {
- std::thread::Builder::new()
+ if let Err(e) = std::thread::Builder::new()
.name("epd-rdma-reaper".into())
.spawn(|| loop {
...
})
- .ok();
+ {
+ error!(error = ?e, "EPD RDMA: failed to spawn reaper thread; slot pool will leak without recycling");
+ }
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| fn spawn_reaper() { | |
| std::thread::Builder::new() | |
| .name("epd-rdma-reaper".into()) | |
| .spawn(|| loop { | |
| std::thread::sleep(REAPER_TICK); | |
| let (Some(g), Some(a)) = ( | |
| AGENT.get().and_then(|o| o.as_ref()), | |
| ARENA.get().and_then(|o| o.as_ref()), | |
| ) else { | |
| continue; | |
| }; | |
| if let Ok(mut notifs) = NotificationMap::new() { | |
| { | |
| let guard = g.lock(); | |
| let _ = guard.agent.get_notifications(&mut notifs, None); | |
| } | |
| if let Ok(map) = notifs.take_notifs() { | |
| for (_agent, tags) in map { | |
| for tag in tags { | |
| if let Ok(room) = tag.parse::<i64>() { | |
| a.pool.free_room(room); | |
| } | |
| } | |
| } | |
| } | |
| } | |
| // TTL sweep: reclaim slots whose READ-notif never arrived. `slot_ttl()` is | |
| // derived to exceed the worker's max hold, so this never races a live READ. | |
| let _ = a.pool.reap_stale(Instant::now(), slot_ttl()); | |
| }) | |
| .ok(); | |
| } | |
| fn spawn_reaper() { | |
| if let Err(e) = std::thread::Builder::new() | |
| .name("epd-rdma-reaper".into()) | |
| .spawn(|| loop { | |
| std::thread::sleep(REAPER_TICK); | |
| let (Some(g), Some(a)) = ( | |
| AGENT.get().and_then(|o| o.as_ref()), | |
| ARENA.get().and_then(|o| o.as_ref()), | |
| ) else { | |
| continue; | |
| }; | |
| if let Ok(mut notifs) = NotificationMap::new() { | |
| { | |
| let guard = g.lock(); | |
| let _ = guard.agent.get_notifications(&mut notifs, None); | |
| } | |
| if let Ok(map) = notifs.take_notifs() { | |
| for (_agent, tags) in map { | |
| for tag in tags { | |
| if let Ok(room) = tag.parse::<i64>() { | |
| a.pool.free_room(room); | |
| } | |
| } | |
| } | |
| } | |
| } | |
| // TTL sweep: reclaim slots whose READ-notif never arrived. `slot_ttl()` is | |
| // derived to exceed the worker's max hold, so this never races a live READ. | |
| let _ = a.pool.reap_stale(Instant::now(), slot_ttl()); | |
| }) { | |
| error!(error = ?e, "EPD RDMA: failed to spawn reaper thread; slot pool will leak without recycling"); | |
| } | |
| } |
🤖 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/mm_rdma/nixl.rs` around lines 423 - 454, The
reaper thread spawn failure is being swallowed in spawn_reaper by discarding the
Result from std::thread::Builder::spawn, so change it to handle the error path
explicitly and log the failure with enough context to diagnose why the
epd-rdma-reaper thread was not created. Keep the existing thread setup and loop
intact, but replace the silent .ok() handling with error reporting tied to
spawn_reaper/Builder::spawn so OS resource exhaustion or similar startup
failures are visible in logs.
| /// Pixel-cache image preprocessing for single-image requests. | ||
| async fn preprocess_image_cached( | ||
| cache: Arc<PixelCache>, | ||
| image: &Arc<ImageFrame>, | ||
| registry: Arc<VisionProcessorRegistry>, | ||
| model_id: String, | ||
| model_type: Option<String>, | ||
| pp_config: PreProcessorConfig, | ||
| fingerprint: u64, | ||
| ) -> Result<PreprocessedEncoderInputs> { | ||
| let key = PixelCacheKey { | ||
| image_hash: image.hash.clone(), | ||
| config_fingerprint: fingerprint, | ||
| }; | ||
| if let Some(cached) = cache.get(&key) { | ||
| return Ok(cached.preprocessed.clone()); | ||
| } | ||
|
|
||
| let preprocessed = preprocess_image_batch( | ||
| registry, | ||
| model_id, | ||
| model_type, | ||
| pp_config, | ||
| std::slice::from_ref(image), | ||
| ) | ||
| .await?; | ||
| cache.insert( | ||
| key, | ||
| Arc::new(CachedPreprocessedItem { | ||
| preprocessed: preprocessed.clone(), | ||
| }), | ||
| ); | ||
| Ok(preprocessed) | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy lift
Concurrent misses on the same image cause duplicate preprocessing (no request coalescing).
preprocess_image_cached does a plain cache-aside check-then-compute-then-insert. If multiple concurrent requests reference the same uncached image (same image_hash + config_fingerprint) — a plausible scenario for a shared/popular image or bursty retries — each one independently misses the cache and redoes the full CPU-bound vision preprocessing before racing to populate the same cache entry. This doesn't corrupt data (the cache converges to one final Arc), but it defeats part of the performance goal of this cache under concurrent load. Consider an in-flight de-duplication map (e.g., DashMap<PixelCacheKey, Arc<tokio::sync::OnceCell<...>>> or similar single-flight pattern) so concurrent misses for the same key await one computation instead of each running it.
🤖 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/multimodal/process.rs` around lines 426 - 459,
The cache-aside flow in preprocess_image_cached still allows concurrent requests
for the same PixelCacheKey to run preprocess_image_batch multiple times before
insert completes. Add in-flight request coalescing around the existing cache
lookup/compute path so only one task computes per key and the others await the
same result, using a single-flight mechanism keyed by PixelCacheKey (for example
with DashMap plus OnceCell or equivalent). Keep the current cache.get and
cache.insert behavior, but wrap the miss path so duplicate misses for the same
image_hash and config_fingerprint share one preprocessing future.
| if let Some(cached) = cache.get(&key) { | ||
| return Ok(cached.preprocessed.clone()); | ||
| } | ||
|
|
||
| let preprocessed = preprocess_image_batch( | ||
| registry, | ||
| model_id, | ||
| model_type, | ||
| pp_config, | ||
| std::slice::from_ref(image), | ||
| ) | ||
| .await?; | ||
| cache.insert( | ||
| key, | ||
| Arc::new(CachedPreprocessedItem { | ||
| preprocessed: preprocessed.clone(), | ||
| }), | ||
| ); | ||
| Ok(preprocessed) |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🔵 Trivial | 🏗️ Heavy lift
Cache hits and inserts both pay a full tensor deep-copy.
cache.get(&key) returns Arc<CachedPreprocessedItem>, but the function immediately does cached.preprocessed.clone() to satisfy the owned PreprocessedEncoderInputs return type — this clones the entire encoder_input ndarray (potentially tens of MB) on every cache hit, and the insert path does an equivalent clone to populate the cache while keeping the original for the caller. This still saves the CPU-bound resize/normalize/patchify work on a hit, but it re-introduces a large per-request allocation+memcpy that the scratch-pooling work elsewhere in this PR is specifically trying to eliminate. If downstream consumers (PrecomputedMultimodalIntermediate.preprocessed) could be changed to hold Arc<PreprocessedEncoderInputs> instead of an owned value, both the hit-path and insert-path clones could be replaced with a cheap Arc clone.
🤖 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/multimodal/process.rs` around lines 440 - 458,
The cache path in process_image_batch is still cloning the full
PreprocessedEncoderInputs on both cache hits and inserts. Update
CachedPreprocessedItem and the downstream
PrecomputedMultimodalIntermediate.preprocessed ownership model so they store
Arc<PreprocessedEncoderInputs> instead of an owned value, then have
process_image_batch return/clones Arc handles rather than deep-copying
encoder_input on every hit or cache fill.
| async fn preprocess_image_batch( | ||
| registry: Arc<VisionProcessorRegistry>, | ||
| model_id: String, | ||
| model_type: Option<String>, | ||
| pp_config: PreProcessorConfig, | ||
| images: &[Arc<ImageFrame>], | ||
| ) -> Result<PreprocessedEncoderInputs> { | ||
| let raw_images: Vec<image::DynamicImage> = images.iter().map(|f| f.image.clone()).collect(); | ||
| tokio::task::spawn_blocking(move || { | ||
| let processor = registry | ||
| .find(&model_id, model_type.as_deref()) | ||
| .ok_or_else(|| anyhow::anyhow!("No vision processor found for model: {model_id}"))?; | ||
| processor | ||
| .preprocess(&raw_images, &pp_config) | ||
| .map_err(|e| anyhow::anyhow!("Image preprocessing failed: {e}")) | ||
| }) | ||
| .await | ||
| .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))? | ||
| } |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
preprocess_image_batch clones DynamicImage outside spawn_blocking, contradicting the file's own established pattern.
images.iter().map(|f| f.image.clone()) runs synchronously in the async function body, before spawn_blocking is invoked — the expensive pixel-buffer clone therefore executes on the tokio async runtime worker thread. This directly contradicts the pattern a few lines above in the same file (the else branch of process_multimodal_parts), which explicitly clones only the cheap Arcs before spawn_blocking and defers f.image.clone() to inside the blocking closure, with the comment: "Extract DynamicImages inside the blocking closure so the expensive clone happens off the tokio async runtime." Since preprocess_image_batch is the new cache-miss path (exercised on every uncached single-image request), this reintroduces exactly the blocking-executor cost the original code was written to avoid.
🔧 Proposed fix: move the DynamicImage clone inside the blocking closure
async fn preprocess_image_batch(
registry: Arc<VisionProcessorRegistry>,
model_id: String,
model_type: Option<String>,
pp_config: PreProcessorConfig,
images: &[Arc<ImageFrame>],
) -> Result<PreprocessedEncoderInputs> {
- let raw_images: Vec<image::DynamicImage> = images.iter().map(|f| f.image.clone()).collect();
+ let images_for_preprocess: Vec<Arc<ImageFrame>> = images.to_vec(); // cheap Arc refcount bumps
tokio::task::spawn_blocking(move || {
+ // Extract DynamicImages inside the blocking closure so the expensive
+ // clone happens off the tokio async runtime.
+ let raw_images: Vec<image::DynamicImage> = images_for_preprocess
+ .iter()
+ .map(|f| f.image.clone())
+ .collect();
let processor = registry
.find(&model_id, model_type.as_deref())
.ok_or_else(|| anyhow::anyhow!("No vision processor found for model: {model_id}"))?;
processor
.preprocess(&raw_images, &pp_config)
.map_err(|e| anyhow::anyhow!("Image preprocessing failed: {e}"))
})
.await
.map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))?
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| async fn preprocess_image_batch( | |
| registry: Arc<VisionProcessorRegistry>, | |
| model_id: String, | |
| model_type: Option<String>, | |
| pp_config: PreProcessorConfig, | |
| images: &[Arc<ImageFrame>], | |
| ) -> Result<PreprocessedEncoderInputs> { | |
| let raw_images: Vec<image::DynamicImage> = images.iter().map(|f| f.image.clone()).collect(); | |
| tokio::task::spawn_blocking(move || { | |
| let processor = registry | |
| .find(&model_id, model_type.as_deref()) | |
| .ok_or_else(|| anyhow::anyhow!("No vision processor found for model: {model_id}"))?; | |
| processor | |
| .preprocess(&raw_images, &pp_config) | |
| .map_err(|e| anyhow::anyhow!("Image preprocessing failed: {e}")) | |
| }) | |
| .await | |
| .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))? | |
| } | |
| async fn preprocess_image_batch( | |
| registry: Arc<VisionProcessorRegistry>, | |
| model_id: String, | |
| model_type: Option<String>, | |
| pp_config: PreProcessorConfig, | |
| images: &[Arc<ImageFrame>], | |
| ) -> Result<PreprocessedEncoderInputs> { | |
| let images_for_preprocess: Vec<Arc<ImageFrame>> = images.to_vec(); // cheap Arc refcount bumps | |
| tokio::task::spawn_blocking(move || { | |
| // Extract DynamicImages inside the blocking closure so the expensive | |
| // clone happens off the tokio async runtime. | |
| let raw_images: Vec<image::DynamicImage> = images_for_preprocess | |
| .iter() | |
| .map(|f| f.image.clone()) | |
| .collect(); | |
| let processor = registry | |
| .find(&model_id, model_type.as_deref()) | |
| .ok_or_else(|| anyhow::anyhow!("No vision processor found for model: {model_id}"))?; | |
| processor | |
| .preprocess(&raw_images, &pp_config) | |
| .map_err(|e| anyhow::anyhow!("Image preprocessing failed: {e}")) | |
| }) | |
| .await | |
| .map_err(|e| anyhow::anyhow!("Preprocessing task panicked: {e}"))? | |
| } |
🤖 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/multimodal/process.rs` around lines 461 - 479,
`preprocess_image_batch` is cloning each `DynamicImage` before `spawn_blocking`,
so the expensive pixel-buffer copy still runs on the async runtime thread.
Update this function to follow the same pattern used in
`process_multimodal_parts` by only collecting the cheap `Arc<ImageFrame>`
references upfront and performing `f.image.clone()` inside the blocking closure.
Keep the `VisionProcessorRegistry::find` and `processor.preprocess` calls in the
blocking task so the full preprocessing work stays off the tokio worker thread.
There was a problem hiding this comment.
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/transforms.rs (1)
149-169: 🚀 Performance & Scalability | 🟡 Minor | ⚡ Quick winReturn the tensor storage to the scratch pool. None of the current
to_tensor*callers hand theArray3<f32>back withinto_raw_vec_and_offset()+scratch::give_f32, soscratch::take_f32here still behaves like a fresh allocation per image. Either recycle the storage in the caller or remove the pool use from this helper.🤖 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/transforms.rs` around lines 149 - 169, The `build_planar_tensor` path is borrowing from `scratch::take_f32` but the resulting `Array3<f32>` storage is never returned, so the pool is not actually reused. Update the `to_tensor*` callers (or this helper) so the tensor’s backing buffer is recovered with `into_raw_vec_and_offset()` and passed back via `scratch::give_f32`, or remove the pool usage from `build_planar_tensor` if recycling won’t be done. Make sure the fix is applied consistently across the `to_tensor*` flow that calls `build_planar_tensor`.
🤖 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.
Outside diff comments:
In `@crates/multimodal/src/vision/transforms.rs`:
- Around line 149-169: The `build_planar_tensor` path is borrowing from
`scratch::take_f32` but the resulting `Array3<f32>` storage is never returned,
so the pool is not actually reused. Update the `to_tensor*` callers (or this
helper) so the tensor’s backing buffer is recovered with
`into_raw_vec_and_offset()` and passed back via `scratch::give_f32`, or remove
the pool usage from `build_planar_tensor` if recycling won’t be done. Make sure
the fix is applied consistently across the `to_tensor*` flow that calls
`build_planar_tensor`.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: 1f86b56f-08ef-4849-a85f-6cebefb83344
📒 Files selected for processing (5)
bindings/python/src/smg/router_args.pycrates/multimodal/src/vision/processors/kimi_k25.rscrates/multimodal/src/vision/transforms.rsmodel_gateway/src/routers/grpc/mm_rdma/nixl.rsmodel_gateway/src/routers/grpc/proto_wrapper.rs
Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
lightseek-bot
left a comment
lightseek-bot
left a comment
There was a problem hiding this comment.
Let's go! @chenht2022 @slin1237
Summary
TensorDatapayload support and TokenSpeed receiver-side NIXL pull path.Stack
This PR is stacked on top of #1852 (
feat(multimodal): add EPD encode routing).Runtime knobs
Default behavior remains inline gRPC payloads. The RDMA pixel path is opt-in.
--features mm-rdmato compile the gateway-side NIXL implementation.SMG_MM_PIXEL_RDMA=1on both the gateway and the TokenSpeed worker that receives pixel payloads. For EPD, that receiver is the encode worker.SMG_RDMA_LISTEN_IP=<gateway-rdma-ip>on the gateway so workers can fetch NIXL metadata.SMG_RDMA_LISTEN_PORTdefaults to18515.SMG_RDMA_SLOT_BYTESon both sides if a serialized image tensor can exceed the default 32MiB slot.Other optimizations in this PR are automatic:
SMG_MM_PIXEL_CACHE_MBSummary by CodeRabbit
rdmaas a multimodal tensor transport option (with documented requirements and fallback behavior).