Fix DCP block alignment before HMA expansion - #10
Draft
pisceskkk wants to merge 18 commits into
Draft
Conversation
A KV connector moves cache at block granularity, so a DCP interleave smaller than block_size cannot be expressed by a block-level transfer. Normalize it in EngineCore before configs reach the executor subprocesses, where model layers cache the value. Signed-off-by: Nicolo Lucchesi <nlucches@redhat.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Identify remote workers by their physical (pp_rank, pcp_rank, tp_rank) key and derive the DCP coverage rank from it, since several physical workers can map to the same DCP rank and the rank alone cannot address an agent. Worker pairing now requires overlap on two axes: KV head coverage and DCP token slice. With both engines on the same dcp_size the slices partition the sequence identically, so logical block ids still line up by index and prefix-cache trimming is unchanged. Heterogeneous DCP and PCP are rejected at handshake rather than silently transferring the wrong token slices. Signed-off-by: Nicolo Lucchesi <nlucches@redhat.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Assert that a decode worker pairs only with prefill workers covering the same token slice, and that heterogeneous DCP or any PCP layout raises at handshake instead of transferring the wrong blocks. Signed-off-by: Nicolo Lucchesi <nlucches@redhat.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
MagicMock rejects attribute access starting with "assert" as a likely assertion typo, so every test using a mocked TransferTopology broke on the handshake guard. validate_ also matches validate_block_size. Signed-off-by: Nicolo Lucchesi <nlucches@redhat.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
… paths compute_tp_mapping already answers "which remote ranks do I exchange KV with, per KV-cache group", and already receives the TransferTopology that carries dcp_rank/dcp_size. Teaching it about DCP keeps one source of truth: with DCP the TP index is head_shard * dcp_size + dcp_rank, so the mapping is computed over head shards and projected back onto remote TP ranks. At dcp_size == 1 every step is an identity. pull_worker and push_worker go back to the upstream shape: specs built from plan.all_source_ranks with per-group gating, the MLA single-read assertion, the MLA notify-the-rest tail, and the push-side attention replication. Only the agent key widens to (pp, pcp, tp). This also restores the three HMA tests the previous structure left red, and drops the MLA duplicate reads: symmetric DCP now maps each local rank to exactly one remote rank per head shard. Signed-off-by: Nicolo Lucchesi <nlucches@redhat.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Leave push_worker on the PR's structure so the diff to review is just pull_worker plus tp_mapping. Its HMA test stays red exactly as it is on the PR branch. Also drop _done_recving_without_xfer: the upstream read path always builds at least one spec per source rank, and an empty block list is already handled by the send_notif branch in _read_blocks, so nothing populates the set any more. Signed-off-by: Nicolo Lucchesi <nlucches@redhat.com> Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
The remote PP stage is fixed once the remote's stages are registered, so cache it per engine in add_remote_agent and drop the lookup from the read path. Also removes the read path's dependency on self.pp_rank. Signed-off-by: Nicolo Lucchesi <nlucches@redhat.com>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
…yers Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: NickLucche <nicolo.lucchesi@mistral.ai>
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>
Signed-off-by: QiuChunshuo <qiuchunshuo@huawei.com>
NickLucche
force-pushed
the
dcp/nixl-tpmapping
branch
from
August 4, 2026 15:06
c2f3f6b to
e4d32ec
Compare
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.
Purpose
Fix DCP block selection when NIXL transfers use hybrid memory allocation (HMA), prefix caching, or an incomplete final DCP stripe.
The read path previously expanded logical block IDs into kernel physical IDs before calling
_dcp_read_slice. This made the DCP phase calculation count physical HMA blocks instead of logical blocks. In addition, equal-DCP transfers did not skip blocks already cached locally, and heterogeneous DCP slices could retain non-overlapping allocation padding.This PR:
_dcp_read_slicecalculation;_dcp_read_slice;Test Result
git diff --checkand the repository sign-off commit hook also pass.Essential Elements of an Effective PR Description Checklist