Repository navigation
Revert "feat(kimi-k3): support P/D disaggregation (#2154)" - #2315
Merged
ZhangLirong-amd merged 1 commit intoSep 20, 2026
Merged
Conversation
This reverts commit d7f20af. The producer now disables deferred output, but the DPA draft lockstep predicate still uses deferred output to exclude pipeline parallelism. Output-less producer prefill chunks skip draft collectives while idle DP peers enter them, hanging V4 DPA P/D requests. Resolve the test_pd_pp conflict by removing the reverted final-state transfer tests and adapting the later network-address fixture to the restored notification-port interface. Preserve subsequent matched-rail routing and network-address coverage.
Contributor
🏷️ CI GuideRuns automatically on every eligible PR before approval:
Heavy model tests:
|
valarLip
approved these changes
Sep 20, 2026
Contributor
There was a problem hiding this comment.
Copilot review overview
🟡 Changes recommended
Unresolved critical and moderate findings remain across notification binding, transfer completion, offload synchronization, and scheduling behavior.
Get a fresh assessment by requesting another Copilot review.
Review effort: Lite
Findings: 5
Open (7)
What changed in this PR
This pull request reverts #2154 to restore prior P/D scheduling and deferred-output behavior, targeting the V4 DPA8/DSpark chunked-prefill hang.
Changes:
- Removes Kimi-K3 P/D state-transfer, slot-region, and offload enhancements.
- Restores prior scheduler, model-runner, forward-dispatch, and KV-transfer behavior.
- Updates related tests and benchmark configurations.
| File | Change summary | Final review status |
|---|---|---|
tests/test_utils.py |
Removes port-zero parsing coverage. | — |
tests/test_scheduler.py |
Restores prior scheduler expectations. | — |
tests/test_pd_pp.py |
Removes reverted state-transfer tests and adapts fixtures. | — |
tests/test_multi_connector.py |
Updates send/save pairing tests. | — |
tests/test_mtp_deferred_status_queue.py |
Removes producer deferred-output tests. | — |
tests/test_lmcache_offload_connector.py |
Restores lookup defaults and removes reverted offload tests. | — |
tests/test_gdn_state_relocation.py |
Removes state-transfer region tests. | — |
tests/test_forward_mode.py |
Removes graph-shape coverage. | — |
tests/test_disagg_modes.py |
Removes full MTP handoff coverage. | — |
tests/test_dense_offload_connector.py |
Removes prefix-reuse coverage. | — |
atom/utils/forward_context.py |
Removes graph-shape dispatch checks. | — |
atom/utils/__init__.py |
Restores prior ZMQ port parsing. | MODERATE (2 votes): Preserve zero as a valid TCP port. |
atom/model_ops/attentions/kimi_mla_gdn_attn.py |
Removes Kimi hybrid state-region transfer integration. | — |
atom/model_ops/attentions/gdn_attn.py |
Removes state-region helpers. | — |
atom/model_engine/scheduler.py |
Restores prior speculation and P/D scheduling behavior. | CRITICAL (1 vote): Preserve full verify width or ragged first-decode handling and retain regression coverage. MODERATE (1 vote): Keep the first-decode flag assignment and regression coverage. |
atom/model_engine/model_runner.py |
Restores deferred output and forward input handling. | — |
atom/kv_transfer/offload/hybrid/kimi_k3/state_object.py |
Removes the separate state-cache budget. | — |
atom/kv_transfer/offload/hybrid/kimi_k3/connector.py |
Restores shared-pool state offload behavior. | — |
atom/kv_transfer/offload/dense/connector.py |
Removes save-stream fencing. | CRITICAL (3 votes): Restore the current-stream fence before asynchronous saves. |
atom/kv_transfer/offload/config.py |
Restores rank-zero lookup defaults. | — |
atom/kv_transfer/offload/chunked_scheduler.py |
Removes load-floor tracking. | MODERATE (3 votes): Restore load-start assignment so loaded prefixes are published. |
atom/kv_transfer/disaggregation/types.py |
Removes producer and remote-slot metadata. | — |
atom/kv_transfer/disaggregation/multi/multi_connector.py |
Updates send/save completion pairing. | CRITICAL (1 vote): Retain terminal handling for saves registered after send completion. |
atom/kv_transfer/disaggregation/mooncake/mooncake_connector.py |
Restores notification-port and transfer behavior. | CRITICAL (3 votes): Retain bound-socket handoff or synchronously report listener bind failure. CRITICAL (1 vote): Retain the asymmetric-only slot-region guard. |
atom/kv_transfer/disaggregation/factory.py |
Restores default backend handling. | — |
atom/kv_transfer/disaggregation/__init__.py |
Removes the producer predicate export. | — |
.github/workflows/atomesh-benchmark.yaml |
Removes reverted benchmark cases. | — |
.github/benchmark/models_atomesh.yaml |
Reverts Kimi benchmark configuration changes. | — |
💡 Configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| self._notification_port: int | None = None | ||
| self._listener_bound = threading.Event() | ||
| self._listener_error: BaseException | None = None | ||
| self._notification_port = get_open_port() |
Comment on lines
+1534
to
+1537
| if has_slot_data and consumer_dcp_size > 1: | ||
| raise RuntimeError( | ||
| "P/D slot-region transfer does not support asymmetric DCP " | ||
| f"(producer_dcp={self.dcp_size}, consumer_dcp={consumer_dcp_size})" | ||
| ) | ||
| if has_slot_data: | ||
| initial_src_slot = int(prefill_data.get("slot_index", -1)) | ||
| prefill_data = dict(prefill_data) | ||
| prefill_data["slot_index"] = final_src_slot | ||
| prefill_data["swa_block_ids"] = request_data.get( | ||
| "src_swa_block_ids", prefill_data.get("swa_block_ids", []) | ||
| ) | ||
| logger.debug( | ||
| "[PD-STATE-SLOT] transfer req_id=%s transfer_id=%s " | ||
| "initial_src_slot=%d final_src_slot=%d dst_slot=%d", | ||
| req_id, | ||
| transfer_id, | ||
| initial_src_slot, | ||
| final_src_slot, | ||
| int(request_data.get("dst_slot_index", -1)), | ||
| "P/D slot-region transfer does not support a DCP consumer " | ||
| f"(consumer_dcp_size={consumer_dcp_size})" |
Comment on lines
390
to
+391
| del self._sent[key] | ||
| rel_save.update(self._saved.pop(key, set())) |
| "save", | ||
| partial(self._do_save_req, producer_event=save_ready_event), | ||
| req, | ||
| self._guard, "save", self._do_save_req, req |
Comment on lines
+2096
to
+2103
| # These trailing slots are the remote's drafts, awaiting | ||
| # verification by the first local decode -- the same contract | ||
| # as postprocess's placeholders, so record the same width. | ||
| # `preempt()` reads it to decide what to strip, and the remote | ||
| # may have sent fewer than `mtp_k` (or none at all), so the | ||
| # count has to come from what was actually appended. T0 is | ||
| # excluded: it is a real generated token. | ||
| seq.num_placeholder_tokens = len(drafts) |
| # The scheduler needs the physical load floor to publish the | ||
| # restored blocks' hashes before resuming suffix prefill. | ||
| seq.offload_load_start_tokens = hbm | ||
| seq.offload_loaded_tokens = self._claim_after_load(seq, hbm, lmc) |
| # Port 0 is a request for the kernel to assign one, not a missing port, so | ||
| # it has to survive the emptiness check below. | ||
| port = "" if parsed.port is None else str(parsed.port) | ||
| port = str(parsed.port or "") |
valarLip
added a commit
that referenced
this pull request
Sep 20, 2026
Two test files conflicted, and both conflicts belong entirely to #2315, which reverts the Kimi-K3 P/D disaggregation of #2154. Every block in them -- `_MLABase` and the transfer-region cases in `test_gdn_state_relocation.py`, the deferred-decode case in `test_mtp_deferred_status_queue.py` -- was added by #2154, and this branch had only adapted them to `KVTransferTensors.set_block_count` and to `decode_spans`. The revert is the right side to take: the code those tests cover is gone. Both files now match main byte for byte. `set_block_count` itself survives the revert with its callers intact. No construction site passes `num_blocks`, and the case pinning that a backend cannot pass one is still there. Gates on the merged tree. The pre-checks suite, which is what CI runs, is 6346 passed / 652 skipped / 0 errors. The full suite against a GPU is 8092 passed and 4 failed, all four in `test_qwen4_exp_mtp.py` and all four inherited: they pass at the merge base and on this branch before the merge, and fail on main alone. #2311 made the QSA GQA kernel's `kv_splits` adapt to the token count, so a batched call and a per-row call now reduce in different orders -- and that test compares them at rtol=0, with the largest difference one bf16 ULP and the selected tokens identical. Nothing caught it upstream because the case needs a GPU and pre-checks runs without one.
wanzhenchn
added a commit
that referenced
this pull request
Sep 21, 2026
`_dp_draft_lockstep_active` read `is_deferred_out` to mean "PP is off", but 47d77d3 also turns it off on a P/D producer. Under DP attention the producer then skipped the aligning `propose()` on an output-less prefill chunk while an idle peer waited in `DPMetadata.make`'s all_reduce -- the V4 DPA8 + DSpark hang #2315 reverted #2154 over. Ask `is_pipeline_parallel` instead.
wanzhenchn
added a commit
that referenced
this pull request
Sep 22, 2026
`_dp_draft_lockstep_active` read `is_deferred_out` to mean "PP is off", but 47d77d3 also turns it off on a P/D producer. Under DP attention the producer then skipped the aligning `propose()` on an output-less prefill chunk while an idle peer waited in `DPMetadata.make`'s all_reduce -- the V4 DPA8 + DSpark hang #2315 reverted #2154 over. Ask `is_pipeline_parallel` instead.
wanzhenchn
added a commit
that referenced
this pull request
Sep 22, 2026
`_dp_draft_lockstep_active` read `is_deferred_out` to mean "PP is off", but 47d77d3 also turns it off on a P/D producer. Under DP attention the producer then skipped the aligning `propose()` on an output-less prefill chunk while an idle peer waited in `DPMetadata.make`'s all_reduce -- the V4 DPA8 + DSpark hang #2315 reverted #2154 over. Ask `is_pipeline_parallel` instead.
wanzhenchn
added a commit
that referenced
this pull request
Oct 8, 2026
`_dp_draft_lockstep_active` read `is_deferred_out` to mean "PP is off", but 47d77d3 also turns it off on a P/D producer. Under DP attention the producer then skipped the aligning `propose()` on an output-less prefill chunk while an idle peer waited in `DPMetadata.make`'s all_reduce -- the V4 DPA8 + DSpark hang #2315 reverted #2154 over. Ask `is_pipeline_parallel` instead.
wanzhenchn
added a commit
that referenced
this pull request
Oct 8, 2026
`_dp_draft_lockstep_active` read `is_deferred_out` to mean "PP is off", but 47d77d3 also turns it off on a P/D producer. Under DP attention the producer then skipped the aligning `propose()` on an output-less prefill chunk while an idle peer waited in `DPMetadata.make`'s all_reduce -- the V4 DPA8 + DSpark hang #2315 reverted #2154 over. Ask `is_pipeline_parallel` instead.
wanzhenchn
added a commit
that referenced
this pull request
Oct 8, 2026
`_dp_draft_lockstep_active` read `is_deferred_out` to mean "PP is off", but 47d77d3 also turns it off on a P/D producer. Under DP attention the producer then skipped the aligning `propose()` on an output-less prefill chunk while an idle peer waited in `DPMetadata.make`'s all_reduce -- the V4 DPA8 + DSpark hang #2315 reverted #2154 over. Ask `is_pipeline_parallel` instead.
wanzhenchn
added a commit
that referenced
this pull request
Oct 10, 2026
`_dp_draft_lockstep_active` read `is_deferred_out` to mean "PP is off", but 47d77d3 also turns it off on a P/D producer. Under DP attention the producer then skipped the aligning `propose()` on an output-less prefill chunk while an idle peer waited in `DPMetadata.make`'s all_reduce -- the V4 DPA8 + DSpark hang #2315 reverted #2154 over. Ask `is_pipeline_parallel` instead.
Yuechguo
pushed a commit
that referenced
this pull request
Oct 10, 2026
* feat(k3-disagg): send the KDA recurrent state, and refuse the transfers that cannot carry it K3's cache is in two halves and the MLA base class describes one: the 24 full-attention layers live in the paged pool it walks, the 69 KDA layers in the slot-indexed recurrent-state pool `build_kv_cache_tensor` tags `per_request_state=True`. A request's KDA state is not derivable from its MLA blocks, so a transfer carrying only the latter leaves decode running most of the model on a zeroed state. * fix(k3-disagg): fail the transfer that registered state slots but carries none The producer skipped Phase 2 whenever either slot index was negative and returned success. Phase 1 has already moved the block regions by then, so on Kimi-K3 the decode side got all 24 MLA layers and a zeroed recurrent state for the other 69 -- and since that still decodes fluent text, nothing downstream could tell. Raise instead, but only where the two facts contradict each other: slot regions were registered, so there is per-request state to move, and no slot came with the transfer to move it from. Backends with nothing slot-indexed registered keep the skip. * fix(k3-disagg): disable deferred output for prefill producers Keep transferred state at the prompt boundary so the consumer processes the first token exactly once, avoiding duplicate KDA recurrent-state updates. * fix(k3-disagg): transfer final KDA slot after prefix checkpoint moves Use the final recurrent-state slot instead of the admission-time slot. Co-authored-by: Cursor <cursoragent@cursor.com> * fix(k3-disagg): complete DSpark P/D handoff compatibility Co-authored-by: Cursor <cursoragent@cursor.com> * fix(k3-disagg): align P/D decode shapes with scheduled batches - derive MLA query width from per-forward batch metadata - fall back to eager execution when an exact CUDA graph shape is unavailable - return the actual staged decode token range - initialize remote first-decode MTP windows without prompt-token leakage - preserve first-decode metadata without an extra scheduler scan - add P/D handoff and shape regression coverage Co-authored-by: Cursor <cursoragent@cursor.com> * fix(k3-disagg): keep the final prefill chunk's token on P/D producers * fix(mooncake): close the side-channel port race that hung P/D requests * fix(offload): fence async KV saves and publish restored prefixes - Fence async KV saves: record a producer event on the RPC stream and wait before background saves read GPU KV, so dense/K3 offload cannot capture unfinished prefill writes. - Publish restored prefixes: preserve the load start offset so the scheduler publishes restored block hashes before suffix prefill resumes. * fix(k3-disagg): allow KDA slot transfer on symmetric DCP8 P/D * test(pd_pp): fix transfer connector stub for kv ready events * fix(offload): unblock save completions, query all ranks, cap K3 state cache * fix(ci-mesh) update Kimi-K3 pd recipe for CI * refactor(k3-disagg): clear the first-decode flag after the batch snapshot. * fix(ci-mesh) align Kimi-K3 pd CI with SA * fix(k3-disagg): keep P/D producers in the DP draft lockstep `_dp_draft_lockstep_active` read `is_deferred_out` to mean "PP is off", but 47d77d3 also turns it off on a P/D producer. Under DP attention the producer then skipped the aligning `propose()` on an output-less prefill chunk while an idle peer waited in `DPMetadata.make`'s all_reduce -- the V4 DPA8 + DSpark hang #2315 reverted #2154 over. Ask `is_pipeline_parallel` instead. * fix(k3-disagg): limit the producer's deferred-output opt-out to recurrent state Deferred output makes a P/D producer decode once past the prompt boundary, so the consumer applies T0 to a state that already has it. That only breaks a recurrent state; a paged write lands at T0's own position and repeating it is idempotent. * fix(k3-disagg): keep EAGLE successor tokens and stop q-bucket from widening producer decodes * fix(k3-disagg): keep producer decodes at one width across DP ranks A P/D producer never speculates but keeps its drafter, so its DP dummy ran at q=mtp_k+1 while real decodes ran at q=1, mismatching `running_tokens` on unified steps. - `dummy_execution` uses the scheduler's spec width (q=1 on a producer). - Only a producer with a drafter passes `graph_shapes`; elsewhere a missing recording stays a KeyError. `decide` warns once on a miss. - Drop the dead spec-token overwrite in the non-deferred `prepare_input_ids` branch; assert `num_spec_step == 0`. - Tests for the above; fix stubs broken by the rebase. --------- Co-authored-by: Cursor <cursoragent@cursor.com>
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.
V4 P/D with DPA8 and DSpark hangs on the first chunked-prefill requests after #2154. That PR makes
is_deferred_out=Falsefor P/D producers, while_dp_draft_lockstep_active()still uses that flag to exclude pipeline parallelism. A producer rank processing an output-less middle prefill chunk skips draft alignment while its idle DP peers enter draft collectives.Reverts #2154 (
d7f20af98c2f93fa50463c9f5282fbb2b5bb49c5) on current main. Resolves thetest_pd_pp.pyconflict by removing tests for reverted state-transfer behavior, and adapts the later network-address test fixture to the restored notification-port interface. Subsequent matched-rail routing and network-address coverage are retained.This is a full rollback: it also withdraws #2154's Kimi-K3 P/D state-handoff, Mooncake listener binding/readiness fix, async offload save fencing, decode graph-shape checks, and associated benchmark changes. In particular, removing listener readiness and save fencing can reintroduce failures outside Kimi-K3, including other Mooncake P/D and offload users. Those fixes need to be reintroduced with V4 DPA + draft coverage.
Validation:
GPU_ARCHS=gfx950: 631 passed, 2 skipped.git diff --checkpass. Ruff reports five existing findings infactory.py/test_utils.py, reproduced unchanged on the base revision.Full GPU validation of this revert is pending.