Skip to content

[PD] Count transferred state bytes per component - #40164

Open
chromecast56 wants to merge 3 commits into
sgl-project:mainfrom
chromecast56:codex/kv-transfer-state-bytes-upstream
Open

chromecast56 wants to merge 3 commits into
sgl-project:mainfrom
chromecast56:codex/kv-transfer-state-bytes-upstream

Conversation

@chromecast56

@chromecast56 chromecast56 commented Sep 18, 2026 •

Copy link
Copy Markdown
Collaborator

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:

  • It multiplies the total index count across state components by the total slot size across components, so a pool with several components (for example a Mamba slot and an SWA page) picks up cross terms.
  • A DSA tail index is a 6-int descriptor of the live slots in one ring row, but it was counted as 6 whole rows. For one live slot of a 4-tensor tail the metric reported 12672 B against 264 B on the wire.

#33806 overlaps on the cross-term fix; this PR also fixes DSA tail accounting and precomputes the per-component sums.

Modifications

  • CommonKVManager computes per-component item-length sums once in __init__ (replacing the unused flat state_item_lens_sum) and exposes get_state_transfer_bytes(component, indices).
  • Regular components count len(indices) * component_bytes; DSA_TAIL counts (first_n + second_n) * row_bytes / tail_size, matching build_dsa_tail_transfer_blocks.
  • The existing replication factor is applied unchanged. Mooncake, NIXL, Ascend (via Mooncake) and Mori all record through _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 bytes build_dsa_tail_transfer_blocks emits for split, single-slot and empty tails. The DSA tail case fails on the previous commit (9504 != 792).

test_dcp_pack.py's hand-built NixlKVManager gains state_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

@chromecast56

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@chromecast56

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@chromecast56
chromecast56 force-pushed the codex/kv-transfer-state-bytes-upstream branch from 0326fc6 to 1f2ea8f Compare September 22, 2026 23:10
@github-actions github-actions Bot added the documentation Improvements or additions to documentation label Sep 22, 2026
@chromecast56

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@chromecast56
chromecast56 force-pushed the codex/kv-transfer-state-bytes-upstream branch from 1f2ea8f to 28aa39f Compare September 22, 2026 23:13
@chromecast56

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@chromecast56
chromecast56 force-pushed the codex/kv-transfer-state-bytes-upstream branch from 28aa39f to fc4dc4e Compare September 23, 2026 00:54
@chromecast56

Copy link
Copy Markdown
Collaborator Author

/tag-and-rerun-ci

@chromecast56
chromecast56 force-pushed the codex/kv-transfer-state-bytes-upstream branch from fc4dc4e to 027294e Compare September 23, 2026 00:58
chromecast56 and others added 3 commits September 23, 2026 21:41
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
chromecast56 force-pushed the codex/kv-transfer-state-bytes-upstream branch from 027294e to 9ccca50 Compare September 23, 2026 21:57
@kpham-sgl kpham-sgl self-assigned this Sep 25, 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]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: unpack indices and give them clear name

@kpham-sgl

Copy link
Copy Markdown
Collaborator

/rerun-failed-ci

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentation Improvements or additions to documentation memory-pool run-ci CI: run the baseline test suite on this PR

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants