Repository navigation
[PD] Count transferred state bytes per component - #40164
Open
chromecast56 wants to merge 3 commits into
Open
chromecast56 wants to merge 3 commits into
chromecast56 wants to merge 3 commits into
Conversation
chromecast56
requested review from
ByronHsu,
Duyi-Wang,
HaiShaw,
ShangmingCai,
hnyls2002 and
sogalin
as code owners
September 18, 2026 10:21
Collaborator
Author
|
/tag-and-rerun-ci |
chromecast56
force-pushed
the
codex/kv-transfer-state-bytes-upstream
branch
2 times, most recently
from
September 22, 2026 21:54
21e5a92 to
0326fc6
Compare
chromecast56
requested review from
Fridge003,
Qiaolin-Yu,
Ying1123,
alphabetc1,
hanming-lu,
huangtingwei9988,
hzh0425,
ispobock,
merrymercy,
xiezhq-hermann and
yizhang2077
as code owners
September 22, 2026 21:54
This was referenced Sep 22, 2026
Collaborator
Author
|
/tag-and-rerun-ci |
chromecast56
force-pushed
the
codex/kv-transfer-state-bytes-upstream
branch
from
September 22, 2026 23:10
0326fc6 to
1f2ea8f
Compare
Collaborator
Author
|
/tag-and-rerun-ci |
chromecast56
force-pushed
the
codex/kv-transfer-state-bytes-upstream
branch
from
September 22, 2026 23:13
1f2ea8f to
28aa39f
Compare
Collaborator
Author
|
/tag-and-rerun-ci |
chromecast56
force-pushed
the
codex/kv-transfer-state-bytes-upstream
branch
from
September 23, 2026 00:54
28aa39f to
fc4dc4e
Compare
Collaborator
Author
|
/tag-and-rerun-ci |
chromecast56
force-pushed
the
codex/kv-transfer-state-bytes-upstream
branch
from
September 23, 2026 00:58
fc4dc4e to
027294e
Compare
A DSA tail index is a 6-int descriptor of the live slots in one ring row, so len(indices) * row_bytes charged six whole rows instead of the slots build_dsa_tail_transfer_blocks puts on the wire. Precompute per-component item-length sums once and drop the unused flat state_item_lens_sum. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Transfer accounting now reads state_item_lens_sums while recording each send, so the hand-built NixlKVManager in test_dcp_pack needs the field. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
chromecast56
force-pushed
the
codex/kv-transfer-state-bytes-upstream
branch
from
September 23, 2026 21:57
027294e to
9ccca50
Compare
kpham-sgl
approved these changes
Sep 29, 2026
| # only the live slots of the one ring row move. | ||
| if len(indices) == 0: | ||
| return 0 | ||
| return (indices[2] + indices[4]) * row_bytes // indices[5] |
Collaborator
There was a problem hiding this comment.
nit: unpack indices and give them clear name
Collaborator
|
/rerun-failed-ci |
This branch has not been deployed
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.
Related PD/DCP/DFlash PRs: #39743, #40794, #39731, #39749. Each is independent and they can merge in any order.
Motivation
The PD KV transfer metric overstates state bytes in two ways:
#33806 overlaps on the cross-term fix; this PR also fixes DSA tail accounting and precomputes the per-component sums.
Modifications
CommonKVManagercomputes per-component item-length sums once in__init__(replacing the unused flatstate_item_lens_sum) and exposesget_state_transfer_bytes(component, indices).len(indices) * component_bytes;DSA_TAILcounts(first_n + second_n) * row_bytes / tail_size, matchingbuild_dsa_tail_transfer_blocks._record_transfer_indices; the fake backend reports no bytes.Pre-existing inaccuracies out of scope: Mooncake/NIXL still record state bytes on hetero-TP ranks that skip the state send, and the hybrid-MLA replica factor stays pinned at 1.
Tests
test/registered/unit/disaggregation/test_kv_transfer_replica_metric.py(5 CPU cases): mixed component slot sizes, and a DSA tail case that compares the metric against the bytesbuild_dsa_tail_transfer_blocksemits for split, single-slot and empty tails. The DSA tail case fails on the previous commit (9504 != 792).test_dcp_pack.py's hand-builtNixlKVManagergainsstate_item_lens_sums, which send-time accounting now reads.CI States
Latest PR Test (Base): ❌ Run #35925585566
Latest PR Test (Extra): ❌ Run #35925585139
Latest PR Test (AMD ROCm 10): ❌ Run #35925585568