refactor(grpc): unify regular/PD/EPD into one Mode-parameterized router + EncodeStage - #1923
Conversation
…D/EPD Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…y guard Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…vacuous guard) Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…cRouter Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…lot on ProcessingState Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…ed by build() Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…ctors Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
…ants Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
📝 WalkthroughWalkthroughThe gRPC routing refactor introduces explicit Regular, PD, and EPD modes, consolidates endpoint pipeline construction, moves multimodal and encode artifacts into shared processing state, and updates router behavior and tests for mode-specific execution. ChangesgRPC mode and factory wiring
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant RouterFactory
participant GrpcRouter
participant RequestPipeline
participant EncodeStage
participant RequestExecutionStage
RouterFactory->>GrpcRouter: create_router(ctx, mode)
GrpcRouter->>RequestPipeline: build(endpoint, mode, deps)
RequestPipeline->>EncodeStage: execute multimodal EPD stage
EncodeStage->>GrpcRouter: store encode_outputs in ProcessingState
RequestExecutionStage->>GrpcRouter: take encode_outputs.dispatch
RequestExecutionStage->>RequestExecutionStage: execute encode/prefill/decode dispatch
Possibly related PRs
Suggested labels: Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Comment |
github-actions
Bot
added
grpc
![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 refactors the gRPC router implementation by consolidating the separate GrpcPDRouter and GrpcRouter into a single, mode-parameterized GrpcRouter capable of serving Regular, PrefillDecode (PD), and EncodePrefillDecode (EPD) modes. It unifies the request pipelines into a builder pattern based on endpoint and mode, introduces a dedicated EncodeStage for EPD, and cleans up redundant files. The review feedback points out a naming collision in encode.rs between the tracing::error! macro and the crate::routers::error module, suggesting an alias to improve clarity.
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.
| use super::PipelineStage; | ||
| use crate::{ | ||
| routers::{ | ||
| error, |
There was a problem hiding this comment.
There is a naming collision and ambiguity between the tracing::error! macro (imported on line 24) and the crate::routers::error module (imported on line 30). While Rust allows this due to separate namespaces for macros and modules, it is highly confusing for readers since error refers to two completely different concepts in the same file. Consider aliasing the router error module to router_error to improve code clarity. If you apply this alias, please also update the usages of the module (such as error::bad_request on line 82) to router_error::bad_request.
| error, | |
| error as router_error, |
![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: 87454ec7ec
ℹ️ 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".
| .ok_or_else(|| format!("gRPC router: no completion pipeline for mode {mode:?}"))?; | ||
|
|
||
| // Regular-only pipelines; `None` in PD/EPD (which 501 these endpoints). | ||
| let harmony_pipeline = RequestPipeline::build(Endpoint::Harmony, mode, &configured_deps); |
There was a problem hiding this comment.
When the router is constructed with Mode::PrefillDecode, this line still builds a Harmony pipeline because RequestPipeline::build(Endpoint::Harmony, ...) only returns None for EPD. route_chat_impl then uses self.harmony_pipeline.is_some() before HarmonyDetector, so any grpc_pd deployment whose model id/card looks like GPT-OSS is diverted into HarmonyPreparation/HarmonyResponseProcessing; the deleted GrpcPDRouter always sent PD chat through the normal PD chat pipeline. Gate this on mode == Mode::Regular (or make the builder return None for PD) to preserve the previous PD routing behavior.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
|
Tested this PR with TokenSpeed EPD on a single 8xH100 node. The end-to-end multimodal path passed with three dedicated GPUs. Result
The latency numbers are smoke-test observations after initialization, not benchmark results. SMG started in the new unified mode: For the successful requests, Prefill's Mooncake sender and Decode's receiver logged the same rendezvous rooms:
Decode logged both requests as No SMG or TokenSpeed runtime source changes were needed. This PR's router/ Environment notes that were necessary for this H100 setup, but do not appear to be regressions from this PR:
Build and run commandsBuild the exact PR head in an isolated checkout: git clone https://github.com/lightseekorg/smg.git /raid/smg-pr1923
git -C /raid/smg-pr1923 fetch origin pull/1923/head
git -C /raid/smg-pr1923 checkout --detach 87454ec7ec9d81308ae6695545d1bf231358fdbd
CARGO_TARGET_DIR=/raid/smg-pr1923-target \
cargo build --release -p smg \
--manifest-path /raid/smg-pr1923/Cargo.toml
sha256sum /raid/smg-pr1923-target/release/smgShared values: export TS_IMAGE=codex/tokenspeed-epd:main-03f644e-ready
export MODEL_HOST=/raid/models/Qwen/Qwen3.5-9B
export MODEL_CONTAINER=/models/Qwen3.5-9BEncode on physical GPU 0: docker run -d --name ts-pr1923-encode \
--gpus '"device=0"' \
--network host --ipc host --pid host --privileged \
-e CUDA_VISIBLE_DEVICES=0 \
-e MC_INTRANODE_NVLINK=1 \
-e TOKENSPEED_SKIP_GRPC_WARMUP=1 \
-v "${MODEL_HOST}:${MODEL_CONTAINER}:ro" \
"${TS_IMAGE}" \
python3 -m smg_grpc_servicer.tokenspeed \
--model "${MODEL_CONTAINER}" \
--served-model-name Qwen3.5-9B \
--host 0.0.0.0 --port 50104 \
--tensor-parallel-size 1 \
--max-model-len 8192 --max-num-seqs 4 \
--gpu-memory-utilization 0.8 \
--attention-backend fa3 \
--disaggregation-mode encode \
--disaggregation-bootstrap-port 18995 \
--disaggregation-transfer-backend mooncake \
--dist-init-addr 127.0.0.1:25000 \
--skip-server-warmupPrefill on physical GPU 1: docker run -d --name ts-pr1923-prefill \
--gpus '"device=1"' \
--network host --ipc host --pid host --privileged \
-e CUDA_VISIBLE_DEVICES=1 \
-e MC_INTRANODE_NVLINK=1 \
-e TOKENSPEED_SKIP_GRPC_WARMUP=1 \
-v "${MODEL_HOST}:${MODEL_CONTAINER}:ro" \
"${TS_IMAGE}" \
python3 -m smg_grpc_servicer.tokenspeed \
--model "${MODEL_CONTAINER}" \
--served-model-name Qwen3.5-9B \
--host 0.0.0.0 --port 50101 \
--tensor-parallel-size 1 \
--max-model-len 8192 --max-num-seqs 4 \
--gpu-memory-utilization 0.8 \
--attention-backend fa3 \
--disaggregation-mode prefill \
--disaggregation-bootstrap-port 19311 \
--disaggregation-transfer-backend mooncake \
--dist-init-addr 127.0.0.1:26000 \
--kvstore-ratio 0.5 \
--enable-prefix-caching \
--enforce-eager \
--skip-server-warmupDecode on physical GPU 2: docker run -d --name ts-pr1923-decode \
--gpus '"device=2"' \
--network host --ipc host --pid host --privileged \
-e CUDA_VISIBLE_DEVICES=2 \
-e MC_INTRANODE_NVLINK=1 \
-e TOKENSPEED_SKIP_GRPC_WARMUP=1 \
-v "${MODEL_HOST}:${MODEL_CONTAINER}:ro" \
"${TS_IMAGE}" \
python3 -m smg_grpc_servicer.tokenspeed \
--model "${MODEL_CONTAINER}" \
--served-model-name Qwen3.5-9B \
--host 0.0.0.0 --port 50111 \
--tensor-parallel-size 1 \
--max-model-len 8192 --max-num-seqs 4 \
--gpu-memory-utilization 0.8 \
--attention-backend fa3 \
--disaggregation-mode decode \
--disaggregation-transfer-backend mooncake \
--dist-init-addr 127.0.0.1:32000 \
--enable-prefix-caching \
--skip-server-warmupRun the exact PR binary as the EPD gateway: docker run -d --name ts-pr1923-gateway \
--network host --ipc host \
-v /raid/smg-pr1923-target/release/smg:/opt/smg-pr1923:ro \
-v "${MODEL_HOST}:${MODEL_CONTAINER}:ro" \
"${TS_IMAGE}" \
/opt/smg-pr1923 launch \
--epd-disaggregation \
--encode grpc://127.0.0.1:50104 18995 \
--prefill grpc://127.0.0.1:50101 19311 \
--decode grpc://127.0.0.1:50111 \
--encode-policy consistent_hashing \
--prefill-policy round_robin \
--decode-policy round_robin \
--policy cache_aware \
--multimodal-tensor-transport inline \
--model-path "${MODEL_CONTAINER}" \
--disable-health-check \
--prometheus-port 29080 \
--host 0.0.0.0 --port 12345Wait for all workers and the gateway: for container in ts-pr1923-encode ts-pr1923-prefill ts-pr1923-decode; do
until docker logs "${container}" 2>&1 | grep -q 'health status -> SERVING'; do
sleep 2
done
done
until curl -fsS http://127.0.0.1:12345/v1/models >/dev/null; do
sleep 2
doneExample multimodal request ( curl -sS http://127.0.0.1:12345/v1/chat/completions \
-H 'Content-Type: application/json' \
--data-binary @- <<'JSON'
{
"model": "Qwen3.5-9B",
"messages": [
{
"role": "user",
"content": [
{
"type": "image_url",
"image_url": {
"url": "data:image/png;base64,<BASE64_PNG>"
}
},
{
"type": "text",
"text": "What is the color of the image? Reply with only the color."
}
]
}
],
"temperature": 0,
"max_tokens": 128,
"stream": false
}
JSON |
lightseek-bot
left a comment
lightseek-bot
left a comment
There was a problem hiding this comment.
LGTM
Thanks for the hard work!!
The Linux-gated dispatch SHM test asserted the segment survived dispatch, but dispatch() disarms the item's own Drop and reclaims the segment via the send-path guard when it returns (success or failure). Assert that real, origin-preserved behavior: a dispatch reclaims the segment exactly once. Signed-off-by: Simo Lin <linsimo.mark@gmail.com>
4 participants
Description
Problem
The regular / PD / EPD axis — which disaggregation topology a request uses — was modeled in three places, three different ways:
GrpcRouterfor regular vsGrpcPDRouterfor PD/EPD) plus a stringly-typedrouter_typelabel field.GrpcPDRouterwas a strict subset ofGrpcRouter(fewer pipelines) plus that label.new_regular/new_pd/new_epd/new_messages{,_pd,_epd}/new_completion{,_pd,_epd}/new_harmony{,_pd}/new_embeddings/new_classify). The stage lists were identical across modes for a given endpoint; only the baked-in(WorkerSelectionMode, ExecutionPlanKind, inject_pd_metadata)args differed.WorkerSelectionMode/ExecutionPlanKindconstructor args).On top of that, EPD's distinctive step — encode — was not a pipeline stage at all. It was smeared across
epd_encode.rs(adapters),plan_epd_encodeincommon/stages/helpers.rs, a duplicated block in both the chat and messagesrequest_buildingstages, and anEncodeDispatchPlancarried inside theExecutionPlan::EncodePrefillDecodeenum variant. TheMultimodalIntermediatewas owned byPreparationOutputand move-consumed in request-building, so encode could not be a separate borrowing stage.Solution
Model the axis once:
Mode { Regular, PrefillDecode, EncodePrefillDecode }(grpc/mode.rs) is the single source of truth. It maps totally to the three stage params and the metrics label:worker_selection(),plan_kind(),inject_pd_metadata(),router_type(). Agrpc_mode(cfg)helper derives it once from(ConnectionMode, RoutingMode).GrpcRouter, parameterized byMode.GrpcPDRouteris deleted. The router carries the regular superset of pipelines with the regular-only ones (harmony,embedding,classify,responses,harmony_responses) asOption, built only whenmode == Regular; PD/EPD leave themNoneand thus return the same trait-default501those endpoints returned before.router_type()returnsself.mode.router_type().RequestPipeline::build(endpoint, mode, deps)(grpc/pipeline.rs), replacing the ~9 constructors. It composes the stage list over the valid(endpoint, mode)matrix and returnsNonefor invalid combos (harmony×EPD, embeddings/classify×PD/EPD — preserving today's gaps). It threadsmodeinto the stages that already accept it.EncodeStage(grpc/common/stages/encode.rs), inserted betweenClientAcquisitionandRequestBuildingonly for EPD pipelines.MultimodalIntermediateis re-homed ontoProcessingState, wherePreparationStagewrites it,EncodeStageborrows it (encode payload, with pixels), andRequestBuildingconsumes it (prefill payload, without pixels). Encode results land in a newProcessingState::encode_outputs { bootstrap_info, dispatch };RequestBuildingreads the bootstrap info andRequestExecutiontakes the dispatch. TheExecutionPlan::EncodePrefillDecodevariant is slimmed to{ request }.Behavior is preserved across every mode: metrics label strings (
"grpc"/"grpc_pd"/"grpc_epd"), per-mode retry worker labels,501responses for regular-only endpoints under PD/EPD, the harmony×EPD gap, encode RPC semantics (fire-and-supervise, wire bootstrap rooms), and the SHM/RDMA cleanup guards (which now move with the owning state, so a cancellation betweenEncodeStageand execution still reclaims/dev/shmsegments viaDrop). The only intentional behavior change is one sanctioned bugfix (below).Changes
Create
grpc/mode.rs—Modeenum +worker_selection/plan_kind/inject_pd_metadata/router_type, andgrpc_mode(cfg)derivation.grpc/common/stages/encode.rs—EncodeStageplus the encode plan/dispatch types and adapters relocated fromepd_encode.rs.Modify
grpc/router.rs— addmode; make regular-only pipelines/contextsOption; build viaRequestPipeline::build;router_type()andDebugfrommode; per-mode metric labels viamatch.grpc/pipeline.rs— replace thenew_*constructors withbuild(endpoint, mode)+ anEndpointenum + the validity matrix; insertEncodeStagefor EPD.grpc/context.rs— re-homemultimodal_intermediateontoProcessingState; addencode_outputs(EncodeOutputs { bootstrap_info, dispatch }); slimExecutionPlan::EncodePrefillDecodeto{ request }.grpc/common/stages/helpers.rs— removeplan_epd_encode(moved intoEncodeStage).grpc/common/stages/request_execution.rs— take the dispatch fromencode_outputsinstead of the enum field.grpc/regular/stages/{chat,messages}/request_building.rs— delete the duplicated inline encode block; readencode_outputs/multimodal_intermediatefrom state.grpc/regular/stages/{chat,messages}/preparation.rs— writemultimodal_intermediateintoProcessingState.routers/factory.rs—create_routerandcreate_igw_routersderiveModeand build the unifiedGrpcRouter;create_grpc_pd_router/create_grpc_epd_routerfolded intocreate_grpc_router(ctx, mode).Delete
grpc/pd_router.rs(GrpcPDRouter) — 535 lines.grpc/epd_encode.rs(logic relocated understages/encode.rs) — 298 lines.Sanctioned bugfix (only intentional behavior change). Pre-refactor
GrpcPDRouter::route_completion_implusedself.retry_configdirectly, skipping the per-model retry override every other method applied. Routing PD/EPD completion through the unifiedroute_completionfixes that by construction: the per-model override now applies uniformly for every mode.Net effect. grpc subsystem: 1645 insertions / 1716 deletions (net -71). Split into production vs test code: production code shrank by ~572 lines, while test code grew by ~501 lines (the parity goldens, the EncodeStage SHM lifecycle tests, and new
Mode/router unit tests).Test Plan
build_matches_frozen_goldens(grpc/pipeline.rs) assertsbuild(endpoint, mode)reproduces a hand-transcribed frozen golden — the exact stage-signature sequence (encodingWorkerSelectionMode,ExecutionPlanKind, andinject_pd_metadata) plus the metrics backend label — for every valid(endpoint, mode). The goldens are transcribed fromorigin/main, not derived frombuild, so a flipped param or swapped mode diverges and fails. Invalid combos (harmony×EPD, embeddings/classify×PD/EPD) assertNone.Mode → (WorkerSelectionMode, ExecutionPlanKind, inject_pd_metadata, router_type)matches the design table exactly.grpc/common/stages/encode.rs, Linux-gated, matching the existing multimodal SHM tests). Backs encode jobs with real/dev/shmsegments and proves: dropping the owningProcessingStatebefore dispatch (cancellation / early return) reclaims the segment viaDrop; a dispatched item transfers SHM ownership off the guard (no double-unlink).router_type()returns"grpc"/"grpc_pd"/"grpc_epd";is_pd_mode()still counts"pd" | "grpc_pd"(EPD excluded); regular-only endpoints return501under PD/EPD; IGW builds all three router instances under the same RouterIds.cargo build -p smg,cargo test -p smg(lib: 1193 passed / 5 ignored; every integration binary — api_tests, routing_tests, spec_test, security_tests, reliability_tests, etc. — 0 failed),cargo clippy -p smg --all-targets,cargo +nightly fmt --check.Checklist
cargo +nightly fmtpassescargo clippy --all-targets --all-features -- -D warningspassesSummary by CodeRabbit
New Features
Bug Fixes
Tests