Repository navigation
feat(disaggregation): add EPD encode pipeline - #548
Conversation
Add embedding disaggregation runtime support for encode workers, transfer, scheduling, and prefill receive. Restructure disaggregation helpers into base/kv/embedding modules and add runtime coverage for EPD encode transfers. Signed-off-by: chenht2022 <chenht2022@gmail.com>
There was a problem hiding this comment.
Pull request overview
This PR adds runtime support for EPD (encode–prefill–decode) multimodal disaggregation by introducing an encode-only worker loop that runs the vision tower, transfers embeddings to prefill over Mooncake, and refactors existing PD KV-disaggregation plumbing into a shared runtime.disaggregation layout.
Changes:
- Add encode-role runtime loop + worker/scheduler/executor, plus a vision-embedding cache (single-tier + optional tiered L1/L2).
- Refactor PD KV-disaggregation into
tokenspeed.runtime.disaggregation.kvwith shared base transport primitives (TransferPoll, bootstrap server, manager base). - Wire EPD admission on prefill (async embedding receive before scheduler admission) and add unit tests for encode scheduling/transfer/cache correctness.
Reviewed changes
Copilot reviewed 53 out of 56 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| tokenspeed-scheduler/csrc/fsm/forward_events.cpp | Clamp finish-time token container to owned pages to avoid over-retract in overlap/disagg paths. |
| test/runtime/test_kv_events.py | Update KV event imports to the new disaggregation KV package. |
| test/runtime/distributed/test_pd_transfer_plan.py | Update PD transfer-plan test imports to new disaggregation KV paths. |
| test/runtime/distributed/test_epd_transfer.py | Add pure-logic tests for encode→prefill embedding wire codec, sharding, and receiver sizing. |
| test/runtime/distributed/test_epd_encode.py | Add tests for encode worker scheduling, executor staging/safety, and cache integration. |
| test/runtime/cache/test_embedding_cache.py | Add unit tests for embedding cache (LRU + tiered L1/L2 behavior). |
| python/tokenspeed/runtime/utils/server_args.py | Add encode disaggregation mode and adjust disaggregation-related arg resolution. |
| python/tokenspeed/runtime/utils/env.py | Add env knob for EPD embedding sharding. |
| python/tokenspeed/runtime/pd/fake/conn.py | Remove legacy PD fake backend implementation (migrated/refactored away). |
| python/tokenspeed/runtime/pd/common/conn.py | Remove legacy PD common backend implementation (migrated/refactored away). |
| python/tokenspeed/runtime/multimodal/inputs.py | Add encode_handshake field to support EPD encode-routed items. |
| python/tokenspeed/runtime/models/qwen3_5.py | Support encoder-only load path; make pre-encode pixel move non-blocking to device. |
| python/tokenspeed/runtime/models/kimi_k25.py | Make pre-encode pixel move non-blocking to device. |
| python/tokenspeed/runtime/models/base/causal_lm.py | Add encoder_only option to skip LM/lm_head/logits allocation for encode role. |
| python/tokenspeed/runtime/layers/attention/backends/base.py | Update StepCounter import to new KV disaggregation path. |
| python/tokenspeed/runtime/execution/model_executor.py | Fix mRoPE fallback logic when only delta is present (decode-only multimodal). |
| python/tokenspeed/runtime/engine/scheduler_utils.py | Disable overlap schedule for encode (encode loop is separate, LM loop not used). |
| python/tokenspeed/runtime/engine/schedule_batch.py | Switch decode scheduler mixin to new disaggregation KV scheduler base class. |
| python/tokenspeed/runtime/engine/request_handler.py | Update BootstrapInfo import to new disaggregation KV types. |
| python/tokenspeed/runtime/engine/event_loop.py | Refactor PD transfer wiring to new KV factory; add EPD prefill admission/drain and encode-loop branch. |
| python/tokenspeed/runtime/engine/async_llm.py | Add encode-request submission path and update disaggregation imports. |
| python/tokenspeed/runtime/disaggregation/utils.py | Add shared DisaggregationMode enum and FastQueue utility. |
| python/tokenspeed/runtime/disaggregation/mooncake_transfer_engine.py | Introduce Mooncake transfer engine wrapper for shared disaggregation usage. |
| python/tokenspeed/runtime/disaggregation/kv/utils.py | Refactor KV disaggregation utilities to use shared TransferPoll and shared mode/queue utilities. |
| python/tokenspeed/runtime/disaggregation/kv/types.py | Define KV disaggregation value types (including BootstrapInfo) under the new package. |
| python/tokenspeed/runtime/disaggregation/kv/transfer_plan.py | Add transfer-plan encoding/decoding and planner logic for PD layouts. |
| python/tokenspeed/runtime/disaggregation/kv/prefill_executor.py | Update KV prefill executor to new TransferPoll and add EPD abort signaling hook. |
| python/tokenspeed/runtime/disaggregation/kv/mooncake/sender.py | Update Mooncake KV sender to use TransferPoll and new module paths. |
| python/tokenspeed/runtime/disaggregation/kv/mooncake/receiver.py | Update Mooncake KV receiver to use TransferPoll and new transfer-plan module. |
| python/tokenspeed/runtime/disaggregation/kv/mooncake/prefill.py | Fix KV-layer indexing for hybrid models; add explicit abort-room notification for EPD failures. |
| python/tokenspeed/runtime/disaggregation/kv/mooncake/entities.py | Update transfer entity imports to new KV transfer-plan/util locations. |
| python/tokenspeed/runtime/disaggregation/kv/mooncake/decode.py | Update decode-side manager to TransferPoll and improve tuple return formatting. |
| python/tokenspeed/runtime/disaggregation/kv/mooncake/conn.py | Add KV Mooncake manager base + bootstrap server based on shared disagg base primitives. |
| python/tokenspeed/runtime/disaggregation/kv/mooncake/async_conn.py | Update async KV path to TransferPoll and new disaggregation utility modules. |
| python/tokenspeed/runtime/disaggregation/kv/mooncake/init.py | Update exports/import paths for refactored Mooncake KV modules. |
| python/tokenspeed/runtime/disaggregation/kv/kv_events.py | Add KV cache event definitions + ZMQ publisher implementation under disaggregation KV. |
| python/tokenspeed/runtime/disaggregation/kv/factory.py | Rename/relocate KV transfer factory (create_kv_transfer) under disaggregation KV. |
| python/tokenspeed/runtime/disaggregation/kv/decode_scheduler.py | Rename scheduler mixin base class for decode scheduling under disaggregation KV. |
| python/tokenspeed/runtime/disaggregation/kv/decode_executor.py | Update decode executor to TransferPoll and prevent repeated FailedEvent emission. |
| python/tokenspeed/runtime/disaggregation/kv/init.py | Replace legacy exports with package docstring describing new module boundaries. |
| python/tokenspeed/runtime/disaggregation/embedding/encode_worker.py | Add encode worker orchestration (cache miss scheduling + cache hit direct ship). |
| python/tokenspeed/runtime/disaggregation/embedding/encode_scheduler.py | Add deterministic token-budget scheduler for vision items. |
| python/tokenspeed/runtime/disaggregation/embedding/encode_loop.py | Add LM-free encode scheduler subprocess loop and cache configuration. |
| python/tokenspeed/runtime/disaggregation/embedding/encode_executor.py | Add ring-buffer staging + safe Mooncake send path for embeddings (incl. deepstack). |
| python/tokenspeed/runtime/disaggregation/embedding/conn.py | Add embedding transport substrate built on shared disaggregation base. |
| python/tokenspeed/runtime/disaggregation/embedding/init.py | Replace legacy exports with package-level documentation for embedding disaggregation. |
| python/tokenspeed/runtime/disaggregation/base/poll.py | Introduce shared TransferPoll status FSM used by KV + embedding transfers. |
| python/tokenspeed/runtime/disaggregation/base/manager.py | Add shared DisaggManagerBase implementing engine + control socket + status tracking. |
| python/tokenspeed/runtime/disaggregation/base/bootstrap.py | Refactor bootstrap rendezvous into role-neutral DisaggBootstrapServer with extension hooks. |
| python/tokenspeed/runtime/disaggregation/base/init.py | Add package docs for shared disaggregation base layer. |
| python/tokenspeed/runtime/configs/model_config.py | Enable encoder-only mode for disaggregation encode and enforce multimodal checkpoint requirement. |
| python/tokenspeed/runtime/cache/prefix_cache.py | Update KV event imports to new disaggregation KV module. |
| python/tokenspeed/runtime/cache/embedding_cache.py | Add embedding cache implementations (LRU + tiered L1/L2). |
| python/tokenspeed/runtime/cache/init.py | Document embedding-cache inclusion in cache subsystem. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| if nbytes > self._ring_bytes: | ||
| raise RuntimeError( | ||
| f"EPD encode embedding {nbytes} B exceeds ring slot " | ||
| f"{self._ring_bytes} B; raise EncodeExecutor._ring_bytes" | ||
| ) |
![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: 844bea3400
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| world_size=server_args.world_size or mapping.world_size, | ||
| dp_size=1, |
There was a problem hiding this comment.
When an encode server is launched with more than one DP worker, mapping.world_size includes all workers but dp_size is hard-coded to 1, so the bootstrap advertises the DP workers as encode TP ranks. The dispatcher still sends each EncodeRequest to just one DP worker (bootstrap_room % len(workers) in data_parallel_controller.py), while prefill receivers fan out across the advertised “TP” ranks; the receivers assigned to other encode workers register rooms those workers never saw, so EPD requests hang or fail. Use the actual encode DP/TP split (or per-DP bootstrap metadata) instead of collapsing DP into TP here.
Useful? React with 👍 / 👎.
| "rank_port": self.rank_port, | ||
| "engine_rank": self.embedding_args.engine_rank, | ||
| } | ||
| url = f"http://{ip}:{self.bootstrap_port}/route" |
There was a problem hiding this comment.
For multi-node encode TP, only attn_tp_rank == 0 starts EmbeddingBootstrapServer, but every rank posts its registration to http://{get_local_ip_by_remote()}:{bootstrap_port}. On non-rank-0 nodes that address is the local host where no bootstrap server was started, so those TP ranks never appear in /route and prefill ranks mapped to them fail to discover an encode endpoint. The registration URL needs to use the rank-0/bootstrap host (for example the distributed init host), not each rank's local IP.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a47f5c0a68
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| for request in new_reqs: | ||
| worker.submit(request) |
There was a problem hiding this comment.
When encode TP has more than one rank and an EncodeRequest carries ShmTensorHandle features, every rank reaches EncodeWorker.submit(), where the handle is attached and then consume() unlinks the POSIX segment. Unlike the normal request path (RequestHandler.recv_reqs() calls sync_shm_features() after broadcast), there is no attach+barrier here, so a faster rank can consume/unlink before a peer has opened the segment, causing that peer's handle.attach() to raise FileNotFoundError and kill the encode loop. This only shows up with SHM-backed encode requests on TP>1.
Useful? React with 👍 / 👎.
| self.executor.register( | ||
| request.request_id, | ||
| request.bootstrap_host, | ||
| request.bootstrap_port, | ||
| request.bootstrap_room, |
There was a problem hiding this comment.
If an EncodeRequest contains more than one multimodal item, all of them are registered under the single request.bootstrap_room here. The embedding transfer protocol is room-keyed and the encode manager marks/pops a room after serving one chunk, while the prefill receiver documents/constructs one receiver per item room; with multiple images sharing this room, later item pre-alloc frames overwrite the earlier room state and only one embedding can be delivered. Either carry per-item rooms or reject multi-item encode requests before scheduling them.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 30506ed376
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if encoder_only: | ||
| # Single model-facing gate: Kimi reads hf_config.encoder_only directly; | ||
| # Qwen3_5ForConditionalGeneration reads it to skip LM construction. | ||
| self.hf_config.encoder_only = True |
There was a problem hiding this comment.
When disaggregation_mode="encode" sets hf_config.encoder_only, Kimi skips constructing language_model, and KimiK25ForConditionalGeneration only casts the vision tower/projector to the LM dtype when self.language_model is not None. For the usual bf16 prefill, the prefill receiver sizes and validates buffers with the prefill model's bf16 vision dtype, while the encoder-only Kimi tower remains float32 and sends str(enc.dtype) == "torch.float32"; validate_fanout_frames() then rejects every transfer with an embedding dtype contract violation. The encode path needs to cast the Kimi vision/projector to the text/prefill dtype even when the LM module is skipped.
Useful? React with 👍 / 👎.
| if self._pause.admit_blocked: | ||
| return |
There was a problem hiding this comment.
When a pause or weight-release request arrives while an EPD request is staged but still waiting for embeddings, this branch stops the EPD drain from admitting or failing it, and the request is not yet in the C++ scheduler. The later PauseController.maybe_finish_drain(self.scheduler) check only observes scheduler state, so it can report the pause drained or release/update memory while a pre-pause EPD request is still pending, then resume and run that request under the new state. Treat _epd_staged/pending EPD receives as in-flight for pause drains or abort/buffer them consistently.
Useful? React with 👍 / 👎.
…lean Signed-off-by: chenht2022 <chenht2022@gmail.com> # Conflicts: # python/tokenspeed/runtime/disaggregation/base/bootstrap.py # python/tokenspeed/runtime/disaggregation/base/poll.py # python/tokenspeed/runtime/disaggregation/embedding/__init__.py # python/tokenspeed/runtime/disaggregation/kv/__init__.py # python/tokenspeed/runtime/disaggregation/kv/mooncake/prefill.py # python/tokenspeed/runtime/disaggregation/kv/types.py # python/tokenspeed/runtime/engine/event_loop.py # python/tokenspeed/runtime/multimodal/inputs.py # python/tokenspeed/runtime/pd/common/conn.py # python/tokenspeed/runtime/pd/fake/conn.py
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 85cb972b58
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| Returns whether the wrapper was installed. | ||
| """ | ||
| if not ( | ||
| hasattr(model, "make_encoder_cudagraph_wrapper") |
There was a problem hiding this comment.
When TOKENSPEED_MM_ENABLE_ENCODER_CUDA_GRAPH is enabled for an EPD encode worker, this guard never passes for the current multimodal models: Qwen3.5 and Kimi expose make_encoder_cudagraph_wrappers (the same plural factory used by ModelExecutor), not make_encoder_cudagraph_wrapper. As a result the encode role silently ignores the CUDA-graph flag and always runs the vision tower eager, which is a significant throughput/latency regression for EPD image/video encoding. Install the wrappers via the plural factory and set the returned encoder attributes as the normal executor does.
Useful? React with 👍 / 👎.
Signed-off-by: chenht2022 <chenht2022@gmail.com>
Signed-off-by: chenht2022 <chenht2022@gmail.com>
…o hongtaoc/epd-encode-clean Signed-off-by: chenht2022 <chenht2022@gmail.com> # Conflicts: # python/pyproject.toml
Signed-off-by: chenht2022 <chenht2022@gmail.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3348835d27
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| else it.recv_main.clone() | ||
| ) | ||
| it.item.encoded_deepstack = None | ||
| it.pool.release(it.pool_slot) |
There was a problem hiding this comment.
With the default pooled receive path for non-deepstack EPD items, clone()/_packed_to_full() on CUDA only enqueues a device copy, but the pool slot is returned immediately. A back-to-back request can lease the same slot and let Mooncake RDMA into it before that copy has read all rows, so the just-published item.encoded can be corrupted under concurrent EPD traffic. Keep the slot leased until a CUDA event/sync confirms the copy out of recv_main has completed.
Useful? React with 👍 / 👎.
| if not self._pending: | ||
| return [], [] # rank-identical emptiness -> all ranks skip the collective |
There was a problem hiding this comment.
For deepstack items, pool-disabled runs, or oversized/exhausted pool fallbacks, _lazy_deregister() schedules the MR/tensor release for the future and then immediately sweeps before the deadline can expire; once that last job is admitted, drain() returns here without polling any job, so an idle prefill worker never revisits _sweep_deregister(). The final fallback receive buffers stay referenced and registered indefinitely until another EPD request arrives, which can strand large GPU buffers after a burst; sweep the deregistration queue even when _pending is empty.
Useful? React with 👍 / 👎.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: b1b0631446
ℹ️ About Codex in GitHub
Your team has set up Codex to 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 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if encoder_only: | ||
| # Single model-facing gate: Kimi reads hf_config.encoder_only directly; | ||
| # Qwen3_5ForConditionalGeneration reads it to skip LM construction. | ||
| self.hf_config.encoder_only = True |
There was a problem hiding this comment.
When this newly enables hf_config.encoder_only, Kimi skips constructing language_model, but KimiK25ForConditionalGeneration.load_weights() still appends every non-vision tensor to language_weights before the final encoder-only gate discards that list. With the safetensors iterator loading CPU tensors shard-by-shard, an encode worker for Kimi still retains the full LM checkpoint in host memory during startup, defeating the vision-only mode and risking OOM/very slow loads; the Kimi loop needs to drop non-vision weights immediately when encoder_only is set.
Useful? React with 👍 / 👎.
| request.bootstrap_room, | ||
| ) | ||
| for idx, item in enumerate(request.items): | ||
| cached = self.cache.get(item.hash) |
There was a problem hiding this comment.
If an encode request carries inline MultimodalDataItems without a prefilled hash (the dataclass default, and submit_encode() does not run the normal input processor's ensure_pad_values() path), every image is looked up and later stored under the None cache key. After the first miss populates cache[None], unrelated images in later encode requests will be treated as cache hits and the wrong embedding will be sent to prefill; call item.ensure_hash() or reject missing hashes before this lookup.
Useful? React with 👍 / 👎.
Summary
runtime.pd, with EPD embedding transfer inruntime.pd.epdand shared bootstrap/manager primitives inruntime.pd.base.How to use
Start TokenSpeed workers in encode/prefill/decode roles:
Launch SMG in EPD mode:
Mooncake is the default disaggregation transfer backend.
Testing
End-to-end EPD was tested with the paired SMG EPD gateway/servicer stack.
Representative setup:
Scaling results: