[Feature] support Chunked Pipeline Parallelism + PD Disaggregation + Atomesh#1552
Draft
Jasen2201 wants to merge 7 commits into
Draft
[Feature] support Chunked Pipeline Parallelism + PD Disaggregation + Atomesh#1552Jasen2201 wants to merge 7 commits into
Jasen2201 wants to merge 7 commits into
Conversation
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
Contributor
There was a problem hiding this comment.
Pull request overview
Adds first-cut pipeline parallelism support (chunked pipeline parallel batching + inter-stage transport) and adjusts KV-cache allocation/indexing to work correctly when each PP stage owns only a slice of layers.
Changes:
- Introduces PP-aware EngineCore topology (one EngineCore per PP stage) plus ZMQ-based metadata/token transport and NCCL-based hidden-state forwarding.
- Updates scheduling/postprocess logic for PP “schedule-time advancement” and in-flight token blocking to keep chunked-prefill correct with multiple batches in flight.
- Fixes KV-cache sizing/keying to allocate per-stage layer counts and index KV tensors by global
layer_numunder PP.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 4 comments.
Show a summary per file
| File | Description |
|---|---|
| tools/pp/test_custom_collective.py | Standalone multi-rank script to sanity-check custom TP collectives under a TP×PP layout. |
| tests/pp/test_pp_transport.py | Unit tests for the new ZMQ PPStageTransport metadata/token round-trip and timeout behavior. |
| tests/pp/test_pp_schedule_advance.py | Unit tests covering schedule-time advancement and postprocess behavior under PP vs pp=1. |
| tests/pp/test_pp_layer_partition.py | Unit tests for PP layer partitioning and make_layers() behavior with missing layers. |
| tests/pp/test_pp_config.py | Unit tests for CLI/config plumbing of pipeline_parallel_size. |
| tests/pp/test_pp_comm.py | Mock-based tests for pp_comm packing/routing of inter-stage tensors and proxy keys. |
| docs/cpp_pp_longctx_repro.md | Repro notes and operational guidance for GLM-5.2 long-context PP bring-up and observed bugs. |
| atom/model_ops/attentions/aiter_mla.py | Uses per-stage total layer count when allocating MLA KV tensors. |
| atom/model_ops/attentions/aiter_attention.py | Uses per-stage total layer count for MHA KV sizing/alloc; aligns memory accounting with PP. |
| atom/model_engine/scheduler.py | Implements PP schedule-time advancement, freezes final-chunk flags, and blocks decode of in-flight tokens. |
| atom/model_engine/pp_engine_core.py | New PP EngineCore implementation: head schedules/collects; downstream executes; last returns tokens. |
| atom/model_engine/model_runner.py | PP-aware device mapping + PP-aware dist init; per-stage layer KV allocation; PP forward path (recv/send intermediates) and deferred-output disabling. |
| atom/model_engine/engine_core.py | Routes run_engine to PP engine core implementation when pipeline_parallel_size > 1. |
| atom/model_engine/engine_core_mgr.py | Spawns one EngineCore per PP stage and routes new requests only to PP heads. |
| atom/model_engine/block_manager.py | Adds start_tokens override for prefix-hash publishing under schedule-time advancement. |
| atom/model_engine/arg_utils.py | Adds pipeline_parallel_size arg + CLI flag -pp/--pipeline-parallel-size. |
| atom/distributed/pp_transport.py | New ZMQ transport layer for PP stage metadata fan-out and token return path. |
| atom/distributed/pp_comm.py | New PP-aware distributed init helper and PP hidden-state/residual send/recv wrappers. |
| atom/config.py | Adds PP-related config/plumbing (pipeline stage rank + ZMQ addr fields + pipeline_parallel_size). |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| results = [] | ||
| for it in range(6): | ||
| # all-reduce: each rank fills with (rig+1); sum should be 1+2+3+4=10 | ||
| x = torch.full((8192,), float(rig + 1), device=rank, dtype=torch.bfloat16) |
| try: | ||
| ag_vals = [] | ||
| for it in range(4): | ||
| x = torch.full((256,), float(rig + 1), device=rank, dtype=torch.bfloat16) |
Comment on lines
+118
to
+120
| fwd_out = self.pp_transport.recv_tokens() | ||
| scheduled_batch, seqs = self._in_flight.popleft() | ||
| self.scheduler.release_pp_inflight(scheduled_batch) |
Comment on lines
+769
to
+771
| """Vestigial: never assigned or read. Worker count is derived directly in | ||
| engine_core from tensor_parallel_size x pipeline_parallel_size x | ||
| prefill_context_parallel_size, not from this field.""" |
zufayu
requested review from
ZhangLirong-amd
and removed request for
ZhangLirong-amd
Add pipeline parallelism support to ATOM using the "one EngineCore per PP stage" architecture (方案②), reusing the existing DP multi-EngineCore infrastructure. Each stage runs as an independent process with its own GPU slice, scheduler (head only), and event loop. Stages communicate hidden states via NCCL P2P and metadata/tokens via direct ZMQ channels. Key changes: - Config: add pipeline_parallel_size (-pp) CLI arg and stage identity - CoreManager: spawn dp*pp EngineCores, route requests to stage-0 heads - PPEngineCoreProc: head/downstream stage loops with lifecycle feedback - pp_comm: PP-aware distributed init (threads pp through aiter groups) - pp_transport: inter-stage ZMQ metadata/token channels - ModelRunner: PP send/recv in run_model, non-last-stage sampling guard, layer_id KV dict key fix (local→global), num_blocks cross-stage reduce-min, disable deferred-output under PP (root cause of decode feeding stale EOS token instead of sampled token to stage-0) Validated on GLM-5.2-MXFP4 (78 layers, 8xMI355X): - tp4 pp2: GSM8K 5-shot = 0.94 (baseline tp8 = 0.931) - tp1 pp8: GSM8K 5-shot = 0.90 (within 100-sample stderr) pp_size=1 path is byte-identical to pre-change behavior. Requires --level 0 --enforce-eager (CUDAGraph+PP deferred to P2).
…ware KV allocation Batch-queue pipeline: the head EngineCore keeps up to pp_size batches in flight instead of blocking on each one (ring=1). The scheduler advances chunked-prefill progress at schedule time (advance_on_schedule) so back-to-back schedule() calls produce successive chunks. Decode seqs whose token is still in the pipeline are blocked from re-scheduling (mark/release_pp_inflight). pp=1 path is byte-identical to phase-1. PP-aware KV allocation: _get_total_num_layers() now returns the local stage's layer count (via get_pp_indices) instead of the full model's num_hidden_layers. All attention builders (MLA, MHA) use this for both compute_block_bytes and allocate_kv_cache_tensors, so each stage only allocates KV for its own layers. Fixes 87% KV waste per stage under PP. Verified: pp2-tp2 GSM8K=0.93 (baseline pp1-tp4=0.95), 121 CPU tests pass. pp8-tp1 has a pre-existing activation-vs-KV-budget issue under investigation.
Prefix the torch profiler output dir with pp{stage}_ when pipeline_parallel_size>1
so each PP stage (all tp-rank 0) writes its trace to its own subdir instead of
colliding in rank_0/. No behavior change at pp=1.
Add docs/cpp_pp_longctx_repro.md: GLM-5.2 PP long-context (60k) prefill bugs
(asm mla_prefill 64-head gate; IndexShare/PP-boundary OOB), fp8 + aligned
VLLM_PP_LAYER_PARTITION workaround, reproduction commands and results.
…ncake push
Enable disaggregated serving with a PP-parallel prefill (e.g. PP4xTP1) feeding
a TP decode (e.g. TP4) via the mooncake push connector.
- port_offset.py: pp-aware side-channel port offset (unique per pp/dp/tp) and
group-major consumer_region_indices mapping (a PP stage's local RDMA regions
-> the consumer's full-layer region list, accounting for the group-major
[kv..., index...] layout aiter_mla registers).
- mooncake_connector.py:
- producer binds a pp-unique side-channel port; scheduler emits remote_pp_size.
- producer writes each region at its group-major consumer index (block + slot
paths); downstream stages (no scheduler) take src_block_ids from the write
request instead of the local prefill cache.
- consumer fans out one write_request per producer stage and requires all
pp_size write-dones before completing a receive.
- consumer-driven release: stage-0 defers reusing the shared page table until
every decode rank confirms all stages wrote (fixes a concurrency KV-corruption
race). pp_size==1 (TP-TP) path is unchanged.
- types.py: ReqMeta.remote_pp_size. proxy.py: forward remote_pp_size /
remote_block_ids to the decode side.
- tests/pp: port offset, group-major mapping (with tiling + regression guards),
per-stage fanout and release completion counting (GPU-free).
Validated on GLM-5.2-MXFP4 (PP4xTP1 + TP4) via mesh: correct generation and
GSM8K parity at concurrency 1; standalone PP4=0.98, TP4=0.97.
…cy accuracy) Two concurrency bugs that corrupted decode KV under load (GSM8K conc=16 0.66): - write-done dedup by stage: the producer resends each write-done 3x for reliability, but the PP consumer counted messages (expected -= 1) instead of distinct stages, so 4 stages x 3 msgs finalized the receive after any 4 arrivals — before lagging stages had written their layers (garbage KV under stage skew). Tag write-done with pp_rank and dedup by distinct stage; finalize only after remote_pp_size distinct stages report. - idempotent _record_release: a duplicate/late release no longer re-adds to done_sending (which would double-free and trip the scheduler's deferred-block assert). Validated GLM-5.2-MXFP4 PP4xTP1 + TP4 via mesh: GSM8K conc=16 0.66 -> 0.96, matching standalone (PP4 0.98 / TP4 0.97).
Under PP the per-stage engine loop does not populate stage-0's _completed_prefills cache, so _wait_for_prefill_data blocked for the full PREFILL_LOOKUP_TIMEOUT (60s) on every request before falling through to the consumer-supplied src_block_ids. Fix: for pp_size>1, all stages (including stage-0) use the consumer- provided src_block_ids directly — never block on the local cache. Non-blocking peek for slot_index only (slot path). pp_size==1 (TP-TP) path unchanged. Result: ttft 60.2s → 0.023s, tpot 83ms → 9.8ms (decode CUDAGraph).
- Extract the duplicated DEALER get-or-create-and-send pattern (start_load_kv, _send_write_done, _send_release) into _send_on_socket(); behavior-preserving (same socket opts, cache, and 3x write-done resend via repeat=3). - Guard cross-thread _release_targets access with _completion_lock for consistency with the rest of the connector state. - Fix stale _pending_recv_slots annotation: dict[ReqId, int] -> tuple[int,int].
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Technical Details
Test Plan
Test Result
Submission Checklist