[PD] Abort prefill KV transfer before page reuse

[AsterWang]

Root cause

In PD mode, decode could abort a request and release its KV pages while prefill-side Mooncake transfers were still in flight. Because prefill had no abort barrier with decode page release, a stale transfer could still land after decode-side reuse and corrupt the new request KV contents.

Test plan

• run PD-separated prefill/decode serving with Mooncake transfer
• force decode-side request timeout / abort on a long-context request
• verify prefill ranks receive abort fanout and return ABORT_ACK
• verify prefill request fails fast instead of hanging after decode abort
• verify decode releases KV pages only after abort ACKs are received
• stress with delayed Mooncake transfer to validate the in-flight write drain path

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[ShangmingCai]

Why? The aborted req's kv indices/slots will be freed later, even if there is still content in flight, this data will not be used, actually. When these indices/slots are reused at the decode side, they will be rewritten with new data, so no stale data as well.

[ShangmingCai]

I assume if the in-flight one is not the last chunk, this request will be aborted by timeout as well, so no stale output. Maybe raise an issue to explore the problem first, then fix it if this is really a potential bug.

[AsterWang]

Why? The aborted req's kv indices/slots will be freed later, even if there is still content in flight, this data will not be used, actually. When these indices/slots are reused at the decode side, they will be rewritten with new data, so no stale data as well.

I agree that the aborted request itself will not use the stale KV slots anymore. The issue we are trying to guard against is a later request reusing the same decode KV slots while the aborted request still has an in-flight remote write targeting those slots.

For example:

T0: old request uses decode KV slot S
T1: prefill has already issued a Mooncake/RDMA write to S, but the write is still in flight
T2: decode times out/aborts the old request and releases slot S
T3: a new request reuses slot S and starts writing/reading its own KV
T4: the old request's in-flight write completes and writes into S

In this ordering, the corrupted output belongs to the new request, not the aborted old request. The new request may rewrite the slot, but without a happens-before relation between the old in-flight write and the slot reuse, the old write can land after the new write and corrupt the new request's KV cache.

[ShangmingCai]

IIUC, T3: a new request reuses slot S and starts writing/reading its own KV will never happen before T4, because prefill batches are executed one by one.

[AsterWang]

IIUC, T3: a new request reuses slot S and starts writing/reading its own KV will never happen before T4, because prefill batches are executed one by one.

I think the "will be rewritten with new data" argument is safe only if we can guarantee the old remote write completes before the new request writes/uses the same slot.

The abort race I am concerned about is exactly the opposite ordering:

T0 old request has an issued Mooncake/RDMA write to slot S
T1 decode aborts old request and releases S
T2 new request reuses S and writes/uses its KV
T3 old request's in-flight write completes after that and overwrites S

If prefill batch serialization guarantees that T3 always happens before T2, then I agree this PR may be unnecessary. My understanding is that Mooncake transfer can remain asynchronous/in-flight after the prefill scheduler has moved on, so batch serialization alone may not provide this transfer-completion barrier.

Do we currently wait for all Mooncake writes of an aborted room to complete before releasing decode KV slots?

[ShangmingCai]

this stale-data case is logically impossible due to the prefill batch-launching logic.

[AsterWang]

this stale-data case is logically impossible due to the prefill batch-launching logic.

https://z.ai/blog/scaling-pain

[ShangmingCai]

this stale-data case is logically impossible due to the prefill batch-launching logic.

https://z.ai/blog/scaling-pain

I think they mean multiple prefill instances writing to the same decode instance, then that could be possible with an extremely small probability (which also requires one super-long prefill with super large chunked prefill size and also requires the decode node to reuse the exact same kv indices for a new incoming req). I think this could be fixed by a small chunked prefill size or a router-side abort for ultra-long input that is impossible to serve within SLO.

In fact, they have proposed some PRs to cover the issue mentioned in the blog, including #23346 and #23539, but no PR for this issue. Could you write a test to make sure we can 100% reproduce this, so that we can verify the fix and ensure no performance drop? Because this is really a rare corner case, we should not sacrifice much for this case unless it is a common case or the fix has zero impact to the performance.

CC: @zRzRzRzRzRzRzR any idea for this?

[ShangmingCai]

IIRC, Mooncake's default transmission timeout is capped at 30 seconds; exceeding this limit triggers a timeout on the Mooncake C++ sync backend. Therefore, the "prefill" step depicted in the diagram within that blog post likely overrides Mooncake's default timeout setting, and the size of the chunked prefill is likely in excess of 200K?

[AsterWang]

IIRC, Mooncake's default transmission timeout is capped at 30 seconds; exceeding this limit triggers a timeout on the Mooncake C++ sync backend. Therefore, the "prefill" step depicted in the diagram within that blog post likely overrides Mooncake's default timeout setting, and the size of the chunked prefill is likely in excess of 200K?

Mooncake timeout protects submitted transfers, but it does not protect stale transfer intents that are still queued on the prefill side. For example, if a chunk is still waiting in the prefill transfer queue when decode hits waiting_timeout and releases the slots, the Mooncake timeout has not started yet. Prefill can later submit that stale write with the old dst_kv_indices.

[ShangmingCai]

@zRzRzRzRzRzRzR do you think this is a proper fix? Will z.ai share the configs, scripts, and a reproduce case for us to verify the fix?

[FrankMinions]

IIRC, Mooncake's default transmission timeout is capped at 30 seconds; exceeding this limit triggers a timeout on the Mooncake C++ sync backend. Therefore, the "prefill" step depicted in the diagram within that blog post likely overrides Mooncake's default timeout setting, and the size of the chunked prefill is likely in excess of 200K?

Mooncake timeout protects submitted transfers, but it does not protect stale transfer intents that are still queued on the prefill side. For example, if a chunk is still waiting in the prefill transfer queue when decode hits waiting_timeout and releases the slots, the Mooncake timeout has not started yet. Prefill can later submit that stale write with the old dst_kv_indices.

I think z.ai may be trying to delay the release of KV slots on the decode side for RDMA writes that failed but are still in-flight, although the probability of this happening should be very low. However, I'm not sure if this is correct; it's just a guess.

[KastanDay]

Hey @AsterWang, @FrankMinions, @ShangmingCai, @zRzRzRzRzRzRzR I opened a complementary PR #32564 that implements the same deferred-release design as yours. And I have a reproducer script (that was mentioned in the above comments here) that shows the failure on current main.
Thanks for your work, I think we've experienced this bug on our servers.

[zRzRzRzRzRzRzR]

Hey @AsterWang, @FrankMinions, @ShangmingCai, @zRzRzRzRzRzRzR I opened a complementary PR #32564 that implements the same deferred-release design as yours. And I have a reproducer script (that was mentioned in the above comments here) that shows the failure on current main. Thanks for your work, I think we've experienced this bug on our servers.

PR #32564 appears to subsume the core abort and deferred-release logic in this PR. Are there any remaining changes here that are not covered by #32564?

In particular, is the additional bounds check for layers_current_pp_stage still necessary? It seems unrelated to the abort fix. If it is needed, could it be split into a separate PR with a regression test? Otherwise, would it make sense to supersede this PR with #32564?

[hnyls2002]

The same root cause (decode aborts while prefill-side transfers are in flight, leading to KV corruption when pages are reused) is addressed by the merged PR #35049 [PD] Deferred decode-side KV release for aborts mid-transfer (2026-08-18): python/sglang/srt/disaggregation/common/conn.py now carries register_deferred_abort_room / note_abort_ack / is_abort_release_safe and the decode side defers page release until every prefill rank has drained its transfer. Closing as fixed - thanks @AsterWang.