diff --git a/python/sglang/srt/disaggregation/base/conn.py b/python/sglang/srt/disaggregation/base/conn.py index 0cc63380b991..6fbc3b990ac1 100644 --- a/python/sglang/srt/disaggregation/base/conn.py +++ b/python/sglang/srt/disaggregation/base/conn.py @@ -132,6 +132,10 @@ def register_to_bootstrap(self): # role switch (release transfer resources; the scheduler owns the KV pool). supports_role_switch: bool = False + # Opt-in per backend: set True once the prefill acks a drained abort + # (ABORT_ACK carrying the sender rank). Without it a hold only ends on timeout. + supports_deferred_decode_kv_release: bool = False + def teardown(self) -> None: raise NotImplementedError( f"{type(self).__name__} does not support PD role switch teardown" diff --git a/python/sglang/srt/disaggregation/common/conn.py b/python/sglang/srt/disaggregation/common/conn.py index 5cf9d060890e..6789c0af8a3a 100644 --- a/python/sglang/srt/disaggregation/common/conn.py +++ b/python/sglang/srt/disaggregation/common/conn.py @@ -182,6 +182,7 @@ def __init__( self.server_args = server_args self.enable_deferred_decode_kv_release = ( envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get() + and self.supports_deferred_decode_kv_release ) self._dcp_pack_buffers = None self._dcp_pack_max_tokens: Optional[int] = None diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py index cdbe95e8089d..b54e876aebbd 100644 --- a/python/sglang/srt/disaggregation/mooncake/conn.py +++ b/python/sglang/srt/disaggregation/mooncake/conn.py @@ -217,6 +217,8 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager): AUX_DATA_HEADER = b"AUX_DATA" # Implements teardown() below, so runtime PD role switching is supported. supports_role_switch = True + # Bootstrap thread defers the ABORT ack until the transfer worker drains. + supports_deferred_decode_kv_release = True def __init__( self, diff --git a/python/sglang/srt/disaggregation/nixl/conn.py b/python/sglang/srt/disaggregation/nixl/conn.py index 8d8b7f963652..cc70acf1d3b0 100644 --- a/python/sglang/srt/disaggregation/nixl/conn.py +++ b/python/sglang/srt/disaggregation/nixl/conn.py @@ -412,6 +412,8 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager): # message is tagged too. It is new to NIXL, hence free to carry the reason. kv_status_msg_tag = b"KV_STATUS" kv_status_msg_carries_reason = True + # ABORT handler defers the ack until the transfer worker drains. + supports_deferred_decode_kv_release = True def __init__( self, diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 4e502e48b9cd..ac02cc073703 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -698,8 +698,9 @@ class Envs: SGLANG_DISAGGREGATION_BOOTSTRAP_ENTRY_CLEANUP_INTERVAL = EnvInt(120) # Deferred decode-side KV release: on abort, hold an in-flight request's KV # pages/slot until the prefill acks the transfer drained, or the timeout - # below fires. Off by default (no behavior/perf impact when disabled). - SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE = EnvBool(False) + # below fires. Only applies to backends that ack the drain + # (supports_deferred_decode_kv_release). + SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE = EnvBool(True) SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE_TIMEOUT = EnvFloat(30.0) # =================================================================== diff --git a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py index 0a8936f219d6..4447f19afe43 100644 --- a/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py +++ b/test/registered/unit/disaggregation/test_deferred_decode_kv_release.py @@ -13,8 +13,12 @@ from unittest.mock import patch from sglang.srt.disaggregation import decode as decode_mod +from sglang.srt.disaggregation.base.conn import BaseKVManager from sglang.srt.disaggregation.common.conn import CommonKVManager from sglang.srt.disaggregation.decode import DecodeTransferQueue +from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager +from sglang.srt.disaggregation.nixl.conn import NixlKVManager +from sglang.srt.environ import envs from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -243,5 +247,24 @@ def test_defer_release_records_deadline_and_idx(self): self.assertIsInstance(deadline, float) +class TestBackendOptIn(CustomTestCase): + """Without a prefill ack, every hold waits out the full release timeout.""" + + def test_enabled_by_default(self): + self.assertTrue(envs.SGLANG_DISAGGREGATION_DEFERRED_DECODE_KV_RELEASE.get()) + + def test_backends_that_ack_opt_in(self): + # Ascend inherits Mooncake's threads, so it opts in too. + for cls in (MooncakeKVManager, NixlKVManager): + with self.subTest(backend=cls.__name__): + self.assertTrue(cls.supports_deferred_decode_kv_release) + + def test_backends_without_a_drain_ack_stay_opted_out(self): + # Inheriting CommonKVManager is not enough: mori marks the room Failed + # without acking. + self.assertFalse(BaseKVManager.supports_deferred_decode_kv_release) + self.assertFalse(CommonKVManager.supports_deferred_decode_kv_release) + + if __name__ == "__main__": unittest.main()