minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[PD] Deferred decode-side KV release for the NIXL backend (#35360)

Browse Source
This commit is contained in:
Shangming Cai
2026-08-19 17:29:45 +08:00
committed by GitHub
parent aa215e5523
commit adca19c497
5 changed files with 308 additions and 65 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/disaggregation/common/conn.py
+44
View File
@@ -236,6 +236,9 @@ class CommonKVManager(BaseKVManager):
)
self.register_to_bootstrap()
self.transfer_infos = {}
# Deferred KV release: aborted room -> (decode_ip, decode_port);
# ack held until the transfer drains.
self._deferred_ack_targets: Dict[int, Tuple[str, int]] = {}
self.req_to_decode_prefix_len: Dict[int, int] = {}
self.decode_kv_args_table = {}
self.pp_group = get_pp_group()
@@ -364,6 +367,47 @@ class CommonKVManager(BaseKVManager):
def clear_deferred_abort_state(self, bootstrap_room: int) -> None:
self._deferred_abort_ack_tracker.pop(bootstrap_room, None)
def _prefill_unique_rank(self) -> int:
"""Stable per-sender id, matching what the transfer worker syncs on Success."""
return (
self.attn_tp_rank * (self.pp_size * self.attn_cp_size)
+ self.pp_rank * self.attn_cp_size
+ self.attn_cp_rank
)
def _send_abort_ack(self, decode_ip: str, decode_port: int, room: int) -> None:
"""Best-effort ack that this rank's transfer for an aborted room drained."""
try:
na = NetworkAddress(decode_ip, decode_port)
self._send_multipart_locked(
na.to_tcp(),
[
b"ABORT_ACK",
str(room).encode("ascii"),
str(self._prefill_unique_rank()).encode("ascii"),
],
is_ipv6=na.is_ipv6,
)
except Exception as e:
logger.debug(f"Failed to send drained ABORT_ACK for room {room}: {e}")
def _maybe_ack_drained_abort(self, room: int) -> None:
"""Send the deferred ack once an aborted room's chunks have drained
(outstanding == 0). pop() makes it fire at most once."""
if self._staging_outstanding.get(room, 0) > 0:
return
target = self._deferred_ack_targets.pop(room, None)
if target is not None:
self._send_abort_ack(target[0], target[1], room)
def register_deferred_ack_target(
self, room: int, decode_ip: str, decode_port: int
) -> None:
"""Hold this room's ack until its transfer drains. Callers must mark the
room Failed FIRST -- registering while it still accepts chunks lets the
worker ack, then a new chunk writes pages the decode already released."""
self._deferred_ack_targets[room] = (decode_ip, decode_port)
def get_kv_replica_factor(self) -> int:
if self._kv_replica_factor is None:
logger.warning_once(
python/sglang/srt/disaggregation/mooncake/conn.py
+18 -48
View File
@@ -8,7 +8,7 @@ import struct
import threading
import time
from collections import defaultdict
from typing import Dict, List, Optional, Set, Tuple, Union
from typing import List, Optional, Set, Tuple, Union
import numpy as np
import numpy.typing as npt
@@ -214,10 +214,6 @@ class MooncakeKVManager(CommonKVManager):
# Per-room count of chunks not yet transferred; teardown waits for
# zero so a deferred chunk is not dropped by an early conclude.
self._staging_outstanding = defaultdict(int)
# Deferred KV release: aborted room -> (decode_ip, decode_port), ack
# held until the transfer drains. Written by the bootstrap thread,
# popped by the single transfer worker that owns the room.
self._deferred_ack_targets: Dict[int, Tuple[str, int]] = {}
self.session_lock = threading.Lock()
# Determine the number of threads to use for kv sender
cpu_count = os.cpu_count()
@@ -1616,39 +1612,6 @@ class MooncakeKVManager(CommonKVManager):
is_ipv6=na.is_ipv6,
)
def _prefill_unique_rank(self) -> int:
"""Stable per-sender id, matching what the transfer worker syncs on Success."""
return (
self.attn_tp_rank * (self.pp_size * self.attn_cp_size)
+ self.pp_rank * self.attn_cp_size
+ self.attn_cp_rank
)
def _send_abort_ack(self, decode_ip: str, decode_port: int, room: int) -> None:
"""Best-effort ack that this rank's transfer for an aborted room drained."""
try:
na = NetworkAddress(decode_ip, decode_port)
self._send_multipart_locked(
na.to_tcp(),
[
b"ABORT_ACK",
str(room).encode("ascii"),
str(self._prefill_unique_rank()).encode("ascii"),
],
is_ipv6=na.is_ipv6,
)
except Exception as e:
logger.debug(f"Failed to send drained ABORT_ACK for room {room}: {e}")
def _maybe_ack_drained_abort(self, room: int) -> None:
"""Send the deferred ack once an aborted room's chunks have drained
(outstanding == 0). pop() makes it fire at most once."""
if self._staging_outstanding.get(room, 0) > 0:
return
target = self._deferred_ack_targets.pop(room, None)
if target is not None:
self._send_abort_ack(target[0], target[1], room)
def transfer_worker(
self,
queue: FastQueue,
@@ -1674,6 +1637,14 @@ class MooncakeKVManager(CommonKVManager):
MooncakeRequestStage.MOONCAKE_WORKER_SEND.level,
)
# Counted at dequeue, before the status check, so
# `outstanding == 0` means nothing is dequeued or in flight --
# the predicate the abort ack relies on. The flag survives
# re-enqueue on defer.
if not kv_chunk.staging_counted:
self._staging_outstanding[kv_chunk.room] += 1
kv_chunk.staging_counted = True
if (
kv_chunk.room not in self.request_status
or self.check_status(kv_chunk.room) == KVPoll.Failed
@@ -1693,11 +1664,6 @@ class MooncakeKVManager(CommonKVManager):
self._maybe_ack_drained_abort(kv_chunk.room)
continue
# Count each chunk once; the flag survives re-enqueue on defer.
if not kv_chunk.staging_counted:
self._staging_outstanding[kv_chunk.room] += 1
kv_chunk.staging_counted = True
if (
self.enable_staging
and staging_strategy is None
@@ -2042,16 +2008,20 @@ class MooncakeKVManager(CommonKVManager):
# flight, decode falls back to the release timeout.
if room_active:
self.update_status(room_to_be_aborted, KVPoll.Failed)
self._deferred_ack_targets[room_to_be_aborted] = (
decode_ip,
decode_port,
self.register_deferred_ack_target(
room_to_be_aborted, decode_ip, decode_port
)
# Try once: the room may already be quiescent and
# never revisited by the worker.
self._maybe_ack_drained_abort(room_to_be_aborted)
logger.debug(
f"Received abort notification for room {room_to_be_aborted}, "
f"marked as Failed; ACK deferred until transfer drains"
)
else:
# Already completed/unknown: no in-flight write, ack now.
elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0:
# Concluded/unknown AND quiescent: ack now. A cleared
# room is not automatically quiescent -- clear() can
# drop a room whose chunk is still transferring.
self._send_abort_ack(
decode_ip, decode_port, room_to_be_aborted
)
python/sglang/srt/disaggregation/nixl/conn.py
+54 -17
View File
@@ -512,7 +512,8 @@ class NixlKVManager(CommonKVManager):
if self.enable_staging:
self._init_staging_decode_ctx()
self._staging_handler = None
self._start_decode_staging_thread()
if self.enable_staging or self.enable_deferred_decode_kv_release:
self._start_decode_listener_thread()
self._start_heartbeat_checker_thread()
else:
raise ValueError(
@@ -592,21 +593,30 @@ class NixlKVManager(CommonKVManager):
return is_watermark_ready(self._staging_ctx, agent_name, alloc_round, alloc_end)
def _start_decode_staging_thread(self):
"""Start a thread on the decode side to recv STAGING_REQ from prefill via ZMQ."""
def _start_decode_listener_thread(self):
"""Decode-side ZMQ listener for STAGING_REQ and ABORT_ACK. A thread, not
NIXL notifs: the decode agent has no progress thread, so notifs only drain
inside a live receiver's poll() and would be missed while idle."""
def decode_staging_thread():
def decode_listener_thread():
while True:
msg = self.server_socket.recv_multipart()
if msg[0] == b"STAGING_REQ":
self._handle_staging_req(msg)
continue
if msg[0] == b"ABORT_ACK":
# Drain ack for an aborted room; aggregate per prefill rank.
if len(msg) >= 3:
self.note_abort_ack(
int(msg[1].decode("ascii")), int(msg[2].decode("ascii"))
)
continue
logger.warning(
"decode_staging_thread: unexpected message tag %s",
"decode_listener_thread: unexpected message tag %s",
msg[0][:20],
)
threading.Thread(target=decode_staging_thread, daemon=True).start()
threading.Thread(target=decode_listener_thread, daemon=True).start()
def _handle_staging_req(self, msg):
from sglang.srt.disaggregation.common.staging_handler import (
@@ -1119,17 +1129,23 @@ class NixlKVManager(CommonKVManager):
room = kv_chunk.room
handles: List[Any] = []
try:
if self.check_status(room) == KVPoll.Failed:
self._staging_outstanding.pop(room, None)
continue
assert room in self.transfer_infos
# Count each chunk once; the flag survives re-enqueue on defer.
# Counted at dequeue, before the status check, so
# `outstanding == 0` means nothing is dequeued or in flight --
# the predicate the abort ack relies on. The flag survives
# re-enqueue on defer.
if not kv_chunk.staging_counted:
self._staging_outstanding[room] += 1
kv_chunk.staging_counted = True
if self.check_status(room) == KVPoll.Failed:
self._staging_outstanding.pop(room, None)
if self.enable_deferred_decode_kv_release:
# Skipped => nothing written for this aborted room; ack.
self._maybe_ack_drained_abort(room)
continue
assert room in self.transfer_infos
# Lazily build a per-worker staging strategy bound to this
# worker's private staging buffer (matches mooncake).
if (
@@ -1344,6 +1360,10 @@ class NixlKVManager(CommonKVManager):
time.sleep(0)
self._staging_outstanding[room] -= 1
if self.enable_deferred_decode_kv_release:
# Handles all DONE => this room's writes landed; ack if it
# was aborted and nothing else is outstanding.
self._maybe_ack_drained_abort(room)
if kv_chunk.is_last_chunk:
self.update_status(room, KVPoll.Success)
elif self.check_status(room) != KVPoll.Success:
@@ -1383,6 +1403,8 @@ class NixlKVManager(CommonKVManager):
self.exceptions[room] = e
self.record_failure(room, str(e))
self.update_status(room, KVPoll.Failed)
# No ack here on purpose: the DONE barrier bails on the first
# ERR, so siblings may still be writing; fall back to the timeout.
def register_buffer_to_engine(self):
self.kv_descs = []
@@ -2628,14 +2650,17 @@ class NixlKVManager(CommonKVManager):
try:
room_to_be_aborted = int(msg[1].decode("ascii"))
decode_ip = msg[2].decode("ascii") if len(msg) > 2 else None
decode_port = int(msg[3].decode("ascii")) if len(msg) > 3 else None
except Exception as e:
logger.debug(f"Ignoring malformed abort notification: {e}")
return True
if (
room_active = (
room_to_be_aborted in self.request_status
and self.check_status(room_to_be_aborted) != KVPoll.Success
):
)
if room_active:
self.record_failure(
room_to_be_aborted,
"Aborted by decode-side abort notification.",
@@ -2651,8 +2676,20 @@ class NixlKVManager(CommonKVManager):
f"ignoring (already completed or unknown)"
)
# TODO: Define real ACK/deferred-release semantics if decode-side buffer
# release needs to wait for prefill-side NIXL transfer quiescence.
# Deferred KV release: register only after the status flip above (see
# register_deferred_ack_target), then try once -- the room may already be
# quiescent and never revisited by the worker. A concluded/unknown room is
# acked only when nothing is still counted for it: the ERR path abandons
# sibling handles that may still be writing and clear() then drops the
# room, so "unknown" alone does not imply quiescent.
if self.enable_deferred_decode_kv_release and decode_port is not None:
if room_active:
self.register_deferred_ack_target(
room_to_be_aborted, decode_ip, decode_port
)
self._maybe_ack_drained_abort(room_to_be_aborted)
elif self._staging_outstanding.get(room_to_be_aborted, 0) == 0:
self._send_abort_ack(decode_ip, decode_port, room_to_be_aborted)
return True