Repository navigation
[PD] Validate decode state layout once at registration - #42051
Merged
ShangmingCai merged 8 commits intoOct 4, 2026
Merged
Conversation
ShangmingCai
requested review from
ByronHsu,
Duyi-Wang,
HaiShaw,
hnyls2002 and
sogalin
as code owners
October 1, 2026 11:08
ShangmingCai
force-pushed
the
pd-state-stride-registration
branch
4 times, most recently
from
October 1, 2026 12:09
11885b6 to
f0d24d1
Compare
Collaborator
Author
|
/tag-and-rerun-ci extra |
github-actions
Bot
added
the
run-ci-extra
Collaborator
Author
|
/rerun-failed-ci |
Collaborator
Author
|
/rerun-group disaggregation |
Contributor
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 ⛔ |
Both inputs to the state stride check are fixed when the peers register, so the verdict is identical for every request on a session. Reach it when the decode peer registers instead of on the first request that happens to carry state indices, and skip the per-transfer comparison for a session already cleared there. A layout rejection gets its own gate rather than joining failed_sessions: that set is probe-recoverable, and an incompatible decode is perfectly reachable, so it would be un-blacklisted straight back into wrong-slot writes. Only a re-registration with a matching layout clears it. _publish_peer_registration records the verdict and publishes decode_kv_args_table last, under one lock, so a transfer can never read a new layout while the previous registration's verdict still stands. _state_stride_mismatch has to repeat _send_kvcache_generic's layout branch so registration can reach it; a test compares the two ASTs so they cannot drift. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
ShangmingCai
force-pushed
the
pd-state-stride-registration
branch
from
October 2, 2026 06:00
f0d24d1 to
d7eb604
Compare
Collaborator
Author
|
/rerun-failed-ci |
Pull the layout branch that _send_kvcache_generic and _state_stride_mismatch both evaluated into one module-level predicate, so the transfer and the registration-time check cannot drift apart. The AST comparison test that pinned the two copies together goes with it. A decode that publishes fewer item lengths than entries used to raise IndexError out of the stride check and take the transfer thread down. Both pairing paths now return it as a rejection reason instead. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Collaborator
Author
|
/rerun-group disaggregation |
Contributor
|
Results for 🚀 🚀 🚀 🚀 🚀 🚀 ⛔ |
The transfer worker now reads state_layout_rejections on every request, so the MooncakeKVManager built via __new__ in test_worker_cleanup_finishes_before_sender_can_report_success needs it alongside failed_sessions. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ender The test still passed dest_tp_ranks/pp_rank, which CommonKVSender.__init__ no longer takes, and its MagicMock manager exposed a truthy deferred_bootstrap, so __init__ took the duplicate-room early exit before reaching the DP-rank branch. Drop the removed args and set deferred_bootstrap=None. Test-only; this fails on main too. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ShangmingCai
added
bypass-maintenance
bypass-fail-fast
Collaborator
Author
|
@aurickq Could you take a quick look of this PR? |
3 of 4 tasks
3 of 4 tasks
ShangmingCai
added a commit
to salexspb/sglang
that referenced
this pull request
Oct 5, 2026
TestMooncakeEarlySend drives the unbound MooncakeKVManager.transfer_worker against a hand-rolled SimpleNamespace. That fake was complete when written, but two changes have since landed on main that made transfer_worker read attributes it does not carry: - sgl-project#42051 added self.state_layout_rejections, read under session_lock for every non-dummy req. - sgl-project#37077 added self.poison_deferred_ack_room(), called from the worker's except path. Merging main into the branch picked up the new conn.py without touching the fake, and the merge was clean because the two sides touch different files. test_worker_waits_before_reading_kv and test_worker_preserves_no_event_path then died with AttributeError on state_layout_rejections, which the except handler turned into a second AttributeError on poison_deferred_ack_room. Only the fixture was stale: the real manager initializes state_layout_rejections on the prefill path in __init__ and inherits poison_deferred_ack_room from CommonKVManager, so no shipping path was affected. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
3 of 5 tasks
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
1 participant
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.
Why is this change needed?
Follow-up to #41607, which validates state entry strides inside
_send_kvcache_genericon every request that carries state.Both inputs to that comparison are fixed when the peers register — the local
kv_args.state_item_lensand the peer'sdst_state_item_lensoffKVArgsRegisterInfo— and neither side changes them for the life of thesession. The only per-request values in that path are the page indices, which
the check never touches, so the verdict is identical for every request on a
session. Two consequences:
carry state indices (
maybe_send_extraruns underif kv_chunk.state_indices and not skip_state), rather than when it registers.PP with compressed MLA that re-runs
_mla_slice_ptrs_for_pp, fourO(num_layers) scans plus list builds, once per state component.
Modifications
_state_stride_mismatch, which returns the firstoffending entry or
None._send_kvcache_genericnow calls it once insteadof carrying two inline copies; the transfer-time behaviour is unchanged (log
the cause,
conclude_failurethe affected room, return -1)._validate_peer_state_layoutwalks every generic state component when adecode peer registers and returns
(fully_checked, mismatch_reason).state_layout_rejections, and the worker'searly exit fails each request to that peer with that exact reason instead of
the generic "Decode instance could be dead" text.
state_strides_validated, and_send_kvcache_genericskips the comparison for it.A layout rejection deliberately does not reuse
failed_sessions. That setis probe-recoverable —
_run_one_probe_passun-blacklists any session whosetransport probe succeeds — and a layout-incompatible decode is perfectly
reachable, so it would be un-blacklisted straight back into wrong-slot writes.
A mismatch is permanent for that peer until it re-registers with a matching
layout, which is also why a session blacklist is the right granularity here
while a transport fault still fails only the affected request.
_publish_peer_registrationrecords the verdict and publishesdecode_kv_args_tablelast, under one lock, so a transfer can never read a newlayout while the previous registration's verdict still stands.
The layout branch that decides between the flat (one region per entry) and
the plain-MHA K/V-split path lives in one module-level predicate,
_uses_flat_entry_layout, which both_send_kvcache_genericand_state_stride_mismatchcall, so the transfer and the registration-time checkcannot drift apart. It is a free function rather than a method because the
test doubles in
test_pp_hybrid_kv_transfer.py/test_mooncake_transfer_batching.pyborrow_send_kvcache_genericwithoutinheriting the class.
A decode that publishes fewer item lengths than entries used to raise
IndexErrorout of the stride check and take the transfer thread down. Bothpairing paths now return that as a rejection reason instead.
Registration does not adjudicate what it cannot see: where the peer published
no sizes for a component, or where entry pairing itself raises, it leaves the
per-transfer check armed so the request still fails with its usual context.
The decode peer is not told at registration time (the register message has no
reply path), so it learns per request, as it does for a transport blacklist.
Test Plan
1a0250c5b6: 9 of the 11 new cases/subtests failthere and all pass with the change. The skip test is driven with entry sizes that would be rejected, so
the transfer going through is proof the comparison was skipped rather than
merely passing.
modules and the two mem_cache modules that also borrow
_send_kvcache_generic: 122 tests and 93 subtests passed. The wholetest/registered/unit/disaggregation/tree shows the same pass/fail counts asthe merged base, plus the new cases.
both directions; that re-registering re-derives it and publishes the layout
last; that a component the peer never published keeps the per-transfer check
armed; that a validated session skips it; that a rejected peer fails every
request with the recorded cause, writes nothing, and never joins the transport
blacklist; that a successful probe clears a transport blacklist but leaves
a layout rejection standing; and that a decode publishing fewer item lengths
than entries yields a rejection reason on both pairing paths (IndexError on
the base).
ruff format,ruff checkandisortclean on both changed files.inspects CPU metadata at handshake and removes work from the transfer path.
🤖 Generated with Claude Code
CI States
Latest PR Test (Base): ⏳ Run #37176704621
Latest PR Test (Extra): 🚫 Run #37184247524
Latest PR Test (AMD ROCm 10): ❌ Run #37176704622