Skip to content

Fix DCP block alignment before HMA expansion - #10

Draft
pisceskkk wants to merge 18 commits into
NickLucche:dcp/nixl-tpmappingfrom
pisceskkk:pr/50611
Draft

Fix DCP block alignment before HMA expansion#10
pisceskkk wants to merge 18 commits into
NickLucche:dcp/nixl-tpmappingfrom
pisceskkk:pr/50611

Conversation

@pisceskkk

@pisceskkk pisceskkk commented Aug 3, 2026

Copy link
Copy Markdown

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:

  • matches local and remote DCP positions using logical block IDs, then expands the matched IDs for HMA;
  • preserves the existing closed-form _dcp_read_slice calculation;
  • applies the prefix offset for equal DCP sizes;
  • truncates local and remote slices to their overlapping block count;
  • clarifies the global-position example in _dcp_read_slice;
  • adds regression coverage for HMA ordering, prefix phase, and local/remote tail padding.

Test Result

9 passed, 14 warnings in 4.19s

git diff --check and the repository sign-off commit hook also pass.


Essential Elements of an Effective PR Description Checklist
  • The purpose and root cause are described.
  • The test plan is provided.
  • The test result is provided.
  • No separate documentation update is required; the affected function documentation is updated inline.

NickLucche and others added 18 commits July 30, 2026 14:21
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants