[Feature] support Chunked Pipeline Parallelism + PD Disaggregation + Atomesh#1552
Draft
Jasen2201 wants to merge 3 commits into
Draft
[Feature] support Chunked Pipeline Parallelism + PD Disaggregation + Atomesh#1552Jasen2201 wants to merge 3 commits into
Jasen2201 wants to merge 3 commits into
Conversation
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.
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.""" |
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