Skip to content

refactor(multimodal): split grpc multimodal.rs into a module directory - #1890

Merged
slin1237 merged 2 commits into
mainfrom
refactor/multimodal-module-split
Jul 8, 2026
Merged

refactor(multimodal): split grpc multimodal.rs into a module directory#1890
slin1237 merged 2 commits into
mainfrom
refactor/multimodal-module-split

Conversation

@slin1237

@slin1237 slin1237 commented Jul 8, 2026

Copy link
Copy Markdown
Member

Description

Problem

model_gateway/src/routers/grpc/multimodal.rs had grown to ~2.65k lines — the
largest flat file under routers/grpc (its neighbors harmony/ and regular/
are already module directories). Everything from content detection to tensor
serialization to SHM transport resolution lived in one file, and the module's
real public surface was buried among ~60 functions with ad-hoc visibility.

Solution

Split multimodal.rs into a multimodal/ module directory along the pipeline's
own phase boundaries. This is a pure move — no behavior change.

Module Responsibility
detect modality detection + content extraction (chat + messages)
config model config-file registry + per-router component bundle
process fetch → preprocess → expand placeholder tokens → build intermediate
assemble intermediate → backend-specific MultimodalData + per-item split
serialize tensor byte/dtype serialization (f32 → bf16/f16)
transport SHM-vs-inline resolution + /dev/shm namespace verification

The crate-facing API is unchanged: the same symbols are re-exported from
mod.rs, so every multimodal::X call site keeps working. Internal helpers are
now enforced-private (pub(super)) instead of implicitly module-scoped, so the
re-export block in mod.rs is now the explicit module contract.

Changes

  • git mv multimodal.rs multimodal/mod.rs; mod.rs now holds the module docs,
    the shared intermediate types (MultimodalOutput, MultimodalIntermediate,
    PrecomputedMultimodalIntermediate), and the pub(crate) re-export surface.
  • Extracted the six submodules above. Tests move next to the code they cover.
  • Net delta +137 lines (per-file imports + module docs); no logic changed.

Reviewing: view with git diff --color-moved=zebra — the vast majority is
verbatim moves. The only non-move edits are per-file use blocks and
visibility bumps (private → pub(super)/pub(crate)).

Test Plan

  • cargo check -p smg --lib --tests — clean
  • cargo clippy -p smg --lib --tests -- -D warnings — clean
  • cargo +nightly fmt — clean (only the new files touched)
  • cargo test -p smg --lib multimodal24 passed, 0 failed (the same
    tests as before, now under assemble/detect/config/process/serialize/transport)

No functional or config surface changed, so no bindings/e2e impact.

Checklist

  • Conventional commit + DCO sign-off
  • cargo +nightly fmt clean
  • cargo clippy -- -D warnings clean
  • Tests pass
  • No behavior change (pure refactor)
  • No config/protocol/bindings surface touched

Summary by CodeRabbit

  • New Features
    • Updated multimodal gRPC handling with image/video detection for chat and Messages API inputs.
    • Added a multimodal request pipeline that preprocesses media, expands multimodal placeholders, and assembles backend-specific payloads.
    • Improved multimodal encoder data transfer, including SHM vs inline selection and efficient tensor serialization with float16/bfloat16 support.
  • Bug Fixes
    • Enhanced multimodal model/preprocessor configuration loading with resilient fallbacks and cached reuse.
    • Added safer SHM lifecycle handling and cleanup during multimodal assembly.

`multimodal.rs` had grown to ~2.65k lines -- the largest flat file under
routers/grpc (harmony/ and regular/ are already directories). Split it by
pipeline phase into a `multimodal/` module directory. This is a pure move
with no behavior change.

- detect:    modality detection + content extraction (chat + messages)
- config:    model config-file registry + per-router component bundle
- process:   fetch -> preprocess -> expand tokens -> build intermediate
- assemble:  intermediate -> backend-specific MultimodalData + per-item split
- serialize: tensor byte/dtype serialization (f32 -> bf16/f16)
- transport: SHM-vs-inline resolution + /dev/shm namespace verification

Tests move next to the code they cover. The crate-facing API is unchanged:
the same symbols are re-exported from mod.rs, and internal helpers are now
enforced-private (pub(super)) instead of implicitly module-scoped.

Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
@github-actions github-actions Bot added grpc gRPC client and router changes model-gateway Model gateway crate changes labels Jul 8, 2026
@coderabbitai

coderabbitai Bot commented Jul 8, 2026

Copy link
Copy Markdown

Review Change Stack

Note

Currently processing new changes in this PR. This may take a few minutes, please wait...

⚙️ Run configuration

Configuration used: Organization UI

Review profile: ASSERTIVE

Plan: Pro

Run ID: 8800f61c-9719-4416-8216-bb9a1c63d924

📥 Commits

Reviewing files that changed from the base of the PR and between c54a0fc and 175f72b.

📒 Files selected for processing (2)
  • model_gateway/src/routers/grpc/multimodal/mod.rs
  • model_gateway/src/routers/grpc/multimodal/process.rs
 ______________________
< Tabs or spaces? Yes. >
 ----------------------
  \
   \   \
        \ /\
        ( )
      .( o ).
📝 Walkthrough

Walkthrough

The monolithic model_gateway/src/routers/grpc/multimodal.rs file (2654 lines) is deleted and its functionality is reimplemented as a new multimodal/ module directory split into mod.rs, detect.rs, config.rs, process.rs, assemble.rs, serialize.rs, and transport.rs, each with corresponding unit tests.

Changes

gRPC multimodal module split

Layer / File(s) Summary
Module wiring and types
model_gateway/src/routers/grpc/multimodal/mod.rs
Declares submodules, re-exports helpers, adds log_mm_timing_enabled, and defines MultimodalOutput, MultimodalIntermediate, PrecomputedMultimodalIntermediate.
Modality detection & extraction
model_gateway/src/routers/grpc/multimodal/detect.rs
Detects image/video modality and extracts MediaContentParts from chat and Messages API inputs, with unit tests.
Config registry
model_gateway/src/routers/grpc/multimodal/config.rs
Adds MultimodalConfigRegistry, MultimodalModelConfig, preprocessor/video config loading with fallbacks, MultimodalComponents wiring, and tests.
Processing pipeline & token expansion
model_gateway/src/routers/grpc/multimodal/process.rs
Resolves placeholder tokens, fetches/preprocesses media on a blocking thread, expands placeholder tokens with structural/patch tracking, returns MultimodalOutput, with tests.
Tensor serialization
model_gateway/src/routers/grpc/multimodal/serialize.rs
Serializes encoder inputs and model-specific tensors with float32/bfloat16/float16 conversion, parallelized 16-bit encoding, and axis-0 slicing, with tests.
SHM transport resolution
model_gateway/src/routers/grpc/multimodal/transport.rs
Resolves encoder dtype and SHM-vs-inline transport based on env/worker overrides and /dev/shm namespace verification, with test.
Backend assembly
model_gateway/src/routers/grpc/multimodal/assemble.rs
Routes intermediate data to Sglang/vLLM/TRT-LLM/TokenSpeed assembly, including SHM cleanup guard, per-item TokenSpeed record construction, and tests.

Estimated code review effort: 4 (Complex) | ~75 minutes

Possibly related issues

Possibly related PRs

  • lightseekorg/smg#495: Original PR that introduced the monolithic multimodal.rs pipeline now split into the new modules.
  • lightseekorg/smg#1515: Implements TokenSpeed proto/servicer support consuming the same assemble_tokenspeed/TokenSpeed multimodal data introduced here.
  • lightseekorg/smg#1604: Overlaps on TokenSpeed SHM tensor transport (shm_namespace_id, SHM cleanup/propagation) directly related to the new transport.rs/assemble.rs.

Suggested labels: multimodal

Suggested reviewers: key4ng, CatherineSue

Poem

A rabbit hopped through tangled code so wide,
One giant file split, now organized inside 🐇
Detect, then config, process, assemble neat,
Transport whispers SHM secrets, oh so sweet,
Tests hop along to keep it all in line! ✨

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title accurately summarizes the main change: refactoring grpc multimodal.rs into a multimodal module directory.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check ✅ Passed Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check ✅ Passed Check skipped because no linked issues were found for this pull request.
✨ Finishing Touches
📝 Generate docstrings
  • Create stacked PR
  • Commit on current branch
🧪 Generate unit tests (beta)
  • Create PR with unit tests
  • Commit unit tests in branch refactor/multimodal-module-split

Comment @coderabbitai help to get the list of available commands.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request refactors the multimodal processing integration by splitting the monolithic multimodal.rs file into a dedicated module structure under model_gateway/src/routers/grpc/multimodal/, dividing the logic into separate files for assembly, configuration, detection, processing, serialization, and transport. The review feedback suggests updating internal documentation links across several of these new files (mod.rs, assemble.rs, and process.rs) to use absolute paths starting with crate:: instead of relative paths, in accordance with internal documentation guidelines.

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.

Comment on lines +6 to +14
//! - [`detect`]: find modalities and extract content parts from chat/messages.
//! - [`config`]: model config-file registry and per-router component bundle.
//! - [`process`]: fetch media → preprocess → expand placeholder tokens →
//! build the lightweight [`MultimodalIntermediate`].
//! - [`assemble`]: turn the intermediate into backend-specific `MultimodalData`
//! once the target backend is known (after worker selection).
//! - [`serialize`]: tensor byte/dtype serialization used by assembly.
//! - [`transport`]: SHM-vs-inline transport resolution and `/dev/shm`
//! namespace verification.

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

According to the repository's internal documentation guidelines, internal documentation links should use absolute paths starting with crate:: instead of relative paths. This ensures that the links remain valid if modules are moved or refactored in the future.

Suggested change
//! - [`detect`]: find modalities and extract content parts from chat/messages.
//! - [`config`]: model config-file registry and per-router component bundle.
//! - [`process`]: fetch media → preprocess → expand placeholder tokens →
//! build the lightweight [`MultimodalIntermediate`].
//! - [`assemble`]: turn the intermediate into backend-specific `MultimodalData`
//! once the target backend is known (after worker selection).
//! - [`serialize`]: tensor byte/dtype serialization used by assembly.
//! - [`transport`]: SHM-vs-inline transport resolution and `/dev/shm`
//! namespace verification.
//! - [`crate::routers::grpc::multimodal::detect`]: find modalities and extract content parts from chat/messages.
//! - [`crate::routers::grpc::multimodal::config`]: model config-file registry and per-router component bundle.
//! - [`crate::routers::grpc::multimodal::process`]: fetch media → preprocess → expand placeholder tokens →
//! build the lightweight [`crate::routers::grpc::multimodal::MultimodalIntermediate`].
//! - [`crate::routers::grpc::multimodal::assemble`]: turn the intermediate into backend-specific `MultimodalData`
//! once the target backend is known (after worker selection).
//! - [`crate::routers::grpc::multimodal::serialize`]: tensor byte/dtype serialization used by assembly.
//! - [`crate::routers::grpc::multimodal::transport`]: SHM-vs-inline transport resolution and `/dev/shm`
//! namespace verification.
References
  1. Use absolute paths starting with crate:: for internal documentation links instead of relative paths (e.g., super::) to ensure links remain valid after refactoring or module moves.

@@ -0,0 +1,782 @@
//! Assembly: convert a [`MultimodalIntermediate`] into backend-specific

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

According to the repository's internal documentation guidelines, internal documentation links should use absolute paths starting with crate:: instead of relative paths. Please update [MultimodalIntermediate] to use its absolute path.

Suggested change
//! Assembly: convert a [`MultimodalIntermediate`] into backend-specific
//! Assembly: convert a [`crate::routers::grpc::multimodal::MultimodalIntermediate`] into backend-specific
References
  1. Use absolute paths starting with crate:: for internal documentation links instead of relative paths (e.g., super::) to ensure links remain valid after refactoring or module moves.

Comment on lines +1 to +5
//! Multimodal processing core: fetch media → preprocess pixels → expand
//! placeholder tokens → build the lightweight [`MultimodalIntermediate`].
//!
//! The chat and Messages API pipelines share `process_multimodal_parts`; only
//! the content extraction differs (see [`super::detect`]).

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

According to the repository's internal documentation guidelines, internal documentation links should use absolute paths starting with crate:: instead of relative paths (such as super::). Please update [MultimodalIntermediate] and [super::detect] to use their absolute paths.

Suggested change
//! Multimodal processing core: fetch media → preprocess pixels → expand
//! placeholder tokens → build the lightweight [`MultimodalIntermediate`].
//!
//! The chat and Messages API pipelines share `process_multimodal_parts`; only
//! the content extraction differs (see [`super::detect`]).
//! Multimodal processing core: fetch media → preprocess pixels → expand
//! placeholder tokens → build the lightweight [`crate::routers::grpc::multimodal::MultimodalIntermediate`].
//!
//! The chat and Messages API pipelines share `process_multimodal_parts`; only
//! the content extraction differs (see [`crate::routers::grpc::multimodal::detect`]).
References
  1. Use absolute paths starting with crate:: for internal documentation links instead of relative paths (e.g., super::) to ensure links remain valid after refactoring or module moves.

@claude claude Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Clean refactoring — verified that all external call sites resolve to the re-exported symbols in mod.rs, visibility modifiers are correct (pub(super) for internal helpers, pub(crate) for the module contract), and no logic was changed. The module split follows the pipeline's natural phase boundaries nicely.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 4

🤖 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 `@model_gateway/src/routers/grpc/multimodal/config.rs`:
- Around line 59-117: The get_or_load flow in MultimodalConfig can trigger
duplicate cold loads for the same tokenizer_id because the cache is checked
before any in-flight load is coordinated. Update the get_or_load path to use a
single-flight or per-tokenizer lock around the resolve_model_config_dir and
config.json/preprocessor_config.json loading so concurrent misses share one
load, then insert the resulting Arc<MultimodalModelConfig> once for that
tokenizer_id.

In `@model_gateway/src/routers/grpc/multimodal/mod.rs`:
- Around line 42-47: The hot-path toggle lookup in log_mm_timing_enabled() is
re-reading and re-parsing SMG_LOG_MM_TIMING on every call from
process_multimodal_parts, which is unnecessary overhead. Cache the parsed
boolean once using a one-time initializer (for example, a OnceLock or
equivalent) inside log_mm_timing_enabled(), and have subsequent calls return the
cached value instead of touching std::env::var repeatedly.

In `@model_gateway/src/routers/grpc/multimodal/process.rs`:
- Around line 427-497: expand_tokens currently warns only when replacements
outnumber placeholder tokens, but it silently leaves extra placeholder_id tokens
in the output once replacement_idx is exhausted. Update the expand_tokens logic
in process.rs to detect and warn on the extra-placeholder case when token_ids
still contain placeholder_id after all PromptReplacement entries are consumed,
while preserving the unchanged token flow. Add a test around expand_tokens
covering the mismatch path so the warning behavior is exercised and the
resulting ExpandedTokens are verified.

In `@model_gateway/src/routers/grpc/multimodal/transport.rs`:
- Around line 149-162: The SHM decision in
transport::worker_matches_shm_namespace is too broad because it returns based on
encode_assignments alone, but shm_enabled is also used for
model_specific_tensors downstream. Update the logic so global SHM is only
enabled when prefill, decode, and all encode workers share the same /dev/shm
namespace, or split the policy into separate flags for encoder_input and
model_specific_tensors. Keep the fix localized in worker_matches_shm_namespace
and the caller path that sets shm_enabled.
🪄 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: 3df68237-1a73-4ed3-9007-b5cfdde5810e

📥 Commits

Reviewing files that changed from the base of the PR and between 5eea1b7 and c54a0fc.

📒 Files selected for processing (8)
  • model_gateway/src/routers/grpc/multimodal.rs
  • model_gateway/src/routers/grpc/multimodal/assemble.rs
  • model_gateway/src/routers/grpc/multimodal/config.rs
  • model_gateway/src/routers/grpc/multimodal/detect.rs
  • model_gateway/src/routers/grpc/multimodal/mod.rs
  • model_gateway/src/routers/grpc/multimodal/process.rs
  • model_gateway/src/routers/grpc/multimodal/serialize.rs
  • model_gateway/src/routers/grpc/multimodal/transport.rs
💤 Files with no reviewable changes (1)
  • model_gateway/src/routers/grpc/multimodal.rs

Comment on lines +59 to +117
pub(crate) async fn get_or_load(
&self,
tokenizer_id: &str,
tokenizer_source: &str,
) -> Result<Arc<MultimodalModelConfig>> {
if let Some(cached) = self.get(tokenizer_id) {
debug!(%tokenizer_id, "multimodal config cache hit");
return Ok(cached);
}

debug!(
%tokenizer_id,
%tokenizer_source,
"multimodal config cache miss, loading"
);

let base_dir = llm_multimodal::hub::resolve_model_config_dir(tokenizer_source)
.await
.with_context(|| {
format!("Failed to resolve model config directory for '{tokenizer_source}'")
})?;

let config_path = base_dir.join("config.json");
let config: serde_json::Value = std::fs::read_to_string(&config_path)
.with_context(|| format!("Failed to read config.json at {}", config_path.display()))
.and_then(|s| {
serde_json::from_str(&s).with_context(|| {
format!("Failed to parse config.json at {}", config_path.display())
})
})?;

// preprocessor_config.json is optional — each vision processor supplies
// its own model-specific defaults, so missing/unparsable files fall
// back to `PreProcessorConfig::default()`. This matches the bundle
// preload path in `try_load_multimodal_config`.
let pp_config_path = base_dir.join("preprocessor_config.json");
let preprocessor_config =
load_preprocessor_config_file(&pp_config_path, "preprocessor_config.json")
.unwrap_or_else(|| {
debug!(
path = %pp_config_path.display(),
"No preprocessor_config.json found; using PreProcessorConfig defaults"
);
PreProcessorConfig::default()
});
let video_preprocessor_config = load_video_preprocessor_config(&base_dir);

let model_config = Arc::new(MultimodalModelConfig {
config,
preprocessor_config,
video_preprocessor_config,
});

self.configs
.insert(tokenizer_id.to_string(), model_config.clone());

debug!(%tokenizer_id, "multimodal config loaded and cached");
Ok(model_config)
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Map relevant files and nearby symbols first.
git ls-files 'model_gateway/src/routers/grpc/multimodal/*' 'llm_multimodal*' | sed -n '1,200p'
echo '---'
ast-grep outline model_gateway/src/routers/grpc/multimodal/config.rs --view expanded
echo '---'
rg -n "resolve_model_config_dir|get_or_load|OnceCell|SingleFlight|preprocessor_config" model_gateway/src llm_multimodal -S

Repository: lightseekorg/smg

Length of output: 9819


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# If the outline is too limited, inspect the implementation of resolve_model_config_dir and the surrounding cache code.
fd -a 'config.rs' model_gateway llm_multimodal . || true
echo '---'
rg -n "pub(crate) async fn get_or_load|struct .*Multimodal|DashMap|resolve_model_config_dir" model_gateway/src -S
echo '---'
rg -n "fn resolve_model_config_dir|async fn resolve_model_config_dir|resolve_model_config_dir" llm_multimodal -S

Repository: lightseekorg/smg

Length of output: 11941


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Look for any code path that indicates network fetching in the model config directory resolver.
rg -n "hub::resolve_model_config_dir|resolve_model_config_dir\(" . -S
echo '---'
rg -n "hf hub|huggingface|network fetch|download|cache miss|snapshot_download|Model::load|hf_hub" . -S

Repository: lightseekorg/smg

Length of output: 24414


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Read only the minimal source slices needed to determine behavior.
echo '--- config.rs ---'
sed -n '1,220p' model_gateway/src/routers/grpc/multimodal/config.rs
echo '--- resolver ---'
for f in $(git ls-files 'llm_multimodal/**' 'llm-multimodal/**' 2>/dev/null); do
  case "$f" in
    *hub*|*config*|*resolve*) echo "### $f"; sed -n '1,260p' "$f"; echo '---';;
  esac
done

Repository: lightseekorg/smg

Length of output: 8122


🏁 Script executed:

#!/bin/bash
set -euo pipefail

sed -n '1,220p' crates/multimodal/src/hub.rs

Repository: lightseekorg/smg

Length of output: 2088


🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Check whether the resolver has any internal deduplication / caching or only plain hf_hub downloads.
rg -n "download|snapshot|cache|OnceCell|Mutex|Semaphore|resolve_model_config_dir" crates/multimodal/src -S

Repository: lightseekorg/smg

Length of output: 1223


Concurrent cold misses can stampede this loader. resolve_model_config_dir downloads config.json and optional processor files for HuggingFace model IDs, so multiple requests for the same uncached tokenizer_id can all perform the same network I/O before the first insert lands. A single-flight cache would collapse the duplicate loads.

🤖 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/config.rs` around lines 59 - 117,
The get_or_load flow in MultimodalConfig can trigger duplicate cold loads for
the same tokenizer_id because the cache is checked before any in-flight load is
coordinated. Update the get_or_load path to use a single-flight or per-tokenizer
lock around the resolve_model_config_dir and
config.json/preprocessor_config.json loading so concurrent misses share one
load, then insert the resulting Arc<MultimodalModelConfig> once for that
tokenizer_id.

Comment thread model_gateway/src/routers/grpc/multimodal/mod.rs
Comment thread model_gateway/src/routers/grpc/multimodal/process.rs
Comment on lines +149 to +162
if !skip_pixel_values {
if let Some(encode_assignments) = encode_assignments {
// EPD: encoder_input (pixels) ships gateway -> encode worker, so SHM
// is safe only if every encode worker assigned in this request shares
// the gateway's /dev/shm. A mixed local/remote fan-out must fall back
// to inline/RDMA rather than giving a remote worker an unreadable SHM handle.
return encode_assignments
.iter()
.all(|assignment| worker_matches_shm_namespace(&assignment.worker, local));
}
}
worker_matches_shm_namespace(prefill, local)
&& worker_matches_shm_namespace(decode, local)
}

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win

Don’t enable global SHM based only on encode workers.

shm_enabled is later used for more than encoder_input: downstream conversion also uses it for model_specific_tensors. Returning after only the encode-worker check can produce SHM-backed tensors that remote prefill/decode workers cannot read. Either require prefill/decode to share /dev/shm too, or split the policy into separate encoder/model-specific SHM flags.

Safer localized fix
-                    return encode_assignments
-                        .iter()
-                        .all(|assignment| worker_matches_shm_namespace(&assignment.worker, local));
+                    let encode_workers_share = encode_assignments
+                        .iter()
+                        .all(|assignment| worker_matches_shm_namespace(&assignment.worker, local));
+                    return encode_workers_share
+                        && worker_matches_shm_namespace(prefill, local)
+                        && worker_matches_shm_namespace(decode, local);
📝 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.

Suggested change
if !skip_pixel_values {
if let Some(encode_assignments) = encode_assignments {
// EPD: encoder_input (pixels) ships gateway -> encode worker, so SHM
// is safe only if every encode worker assigned in this request shares
// the gateway's /dev/shm. A mixed local/remote fan-out must fall back
// to inline/RDMA rather than giving a remote worker an unreadable SHM handle.
return encode_assignments
.iter()
.all(|assignment| worker_matches_shm_namespace(&assignment.worker, local));
}
}
worker_matches_shm_namespace(prefill, local)
&& worker_matches_shm_namespace(decode, local)
}
if !skip_pixel_values {
if let Some(encode_assignments) = encode_assignments {
// EPD: encoder_input (pixels) ships gateway -> encode worker, so SHM
// is safe only if every encode worker assigned in this request shares
// the gateway's /dev/shm. A mixed local/remote fan-out must fall back
// to inline/RDMA rather than giving a remote worker an unreadable SHM handle.
let encode_workers_share = encode_assignments
.iter()
.all(|assignment| worker_matches_shm_namespace(&assignment.worker, local));
return encode_workers_share
&& worker_matches_shm_namespace(prefill, local)
&& worker_matches_shm_namespace(decode, local);
}
}
worker_matches_shm_namespace(prefill, local)
&& worker_matches_shm_namespace(decode, local)
}
🤖 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/transport.rs` around lines 149 -
162, The SHM decision in transport::worker_matches_shm_namespace is too broad
because it returns based on encode_assignments alone, but shm_enabled is also
used for model_specific_tensors downstream. Update the logic so global SHM is
only enabled when prefill, decode, and all encode workers share the same
/dev/shm namespace, or split the policy into separate flags for encoder_input
and model_specific_tensors. Keep the fix localized in
worker_matches_shm_namespace and the caller path that sets shm_enabled.

Two review fixes on top of the module split, kept as a separate commit so the
split itself stays a pure, reviewable move:

- log_mm_timing_enabled: read SMG_LOG_MM_TIMING once via OnceLock instead of
  re-parsing the env on every multimodal request.
- expand_tokens: warn when the token sequence has more placeholder tokens than
  PromptReplacements (the excess were previously left unexpanded silently).
  Token output is unchanged; adds a regression test.

Both address CodeRabbit review comments. The single-flight concern in
MultimodalConfigRegistry::get_or_load is deferred to a dedicated PR (it needs
real async coordination + concurrency tests, out of scope for this refactor).

Signed-off-by: Simo Lin <25425177+slin1237@users.noreply.github.com>
@slin1237

slin1237 commented Jul 8, 2026

Copy link
Copy Markdown
Member Author

Thanks for the reviews. Dispositions:

  • log_mm_timing_enabled env re-read — fixed in 175f72b (cached via OnceLock, read once instead of per request).
  • expand_tokens silent excess placeholders — fixed in 175f72b (added a warning + regression test; token output is unchanged).
  • get_or_load single-flight — valid, but it's a real async-coordination change (per-key single-flight + concurrency tests) and out of scope for a pure-move refactor. The registry is preloaded at tokenizer registration so the cold-miss race is rare; tracking it as a dedicated follow-up.
  • transport SHM "too broad" — this looks like a misread of the EPD data flow: the SHM-backed encoder_input and model_specific_tensors are dispatched to the encode workers, while the prefill/decode legs call clear_mm_pixel_values() and receive the embedding over Mooncake. So gating shm_enabled on encode_assignments sharing the gateway's /dev/shm is correct — there is no second destination that would need a split flag. No change.
  • Gemini doc-link paths — the intra-doc links resolve correctly and the repo has no absolute-path convention (existing code mixes crate::, super::, and bare [Type]); keeping the concise relative form for a module listing its own submodules.

For context: this PR is a pure module split (commit 1) plus the two small review fixes above (commit 2), kept separate so the move stays reviewable as a no-op via git diff --color-moved.

@slin1237
slin1237 merged commit 0cc54c7 into main Jul 8, 2026
13 of 17 checks passed
@slin1237
slin1237 deleted the refactor/multimodal-module-split branch July 8, 2026 18:11
@coderabbitai

coderabbitai Bot commented Jul 8, 2026

Copy link
Copy Markdown

Caution

Failed to replace (edit) comment. This is likely due to insufficient permissions or the comment being deleted.

Error details
putComment timed out

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

grpc gRPC client and router changes model-gateway Model gateway crate changes

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant