[PD] Support fake decode for PD disaggregation without prefill node (#14628)
Co-authored-by: sunhailiang <sunhailiang@baidu.com>
This commit is contained in:
co-authored by
sunhailiang
parent
061f41affc
commit
bc3ca30023
@@ -323,7 +323,11 @@ class DecodePreallocQueue:
|
|||||||
req.retraction_mb_id = None
|
req.retraction_mb_id = None
|
||||||
self.retracted_queue.append(req)
|
self.retracted_queue.append(req)
|
||||||
else:
|
else:
|
||||||
if req.bootstrap_host == FAKE_BOOTSTRAP_HOST:
|
# Auto enable FAKE mode if configured
|
||||||
|
if req.bootstrap_host == FAKE_BOOTSTRAP_HOST or (
|
||||||
|
req.bootstrap_host is None
|
||||||
|
and self.scheduler.server_args.disaggregation_decode_enable_fake_auto
|
||||||
|
):
|
||||||
kv_receiver_class = get_kv_class(
|
kv_receiver_class = get_kv_class(
|
||||||
TransferBackend.FAKE, KVClassType.RECEIVER
|
TransferBackend.FAKE, KVClassType.RECEIVER
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -489,6 +489,16 @@ class DataParallelController:
|
|||||||
self.workers
|
self.workers
|
||||||
)
|
)
|
||||||
else:
|
else:
|
||||||
|
# Set default bootstrap_room if in FAKE auto mode and room is None
|
||||||
|
if (
|
||||||
|
req.bootstrap_room is None
|
||||||
|
and self.server_args.disaggregation_decode_enable_fake_auto
|
||||||
|
):
|
||||||
|
req.bootstrap_room = self.round_robin_counter
|
||||||
|
self.round_robin_counter = (self.round_robin_counter + 1) % len(
|
||||||
|
self.workers
|
||||||
|
)
|
||||||
|
|
||||||
assert (
|
assert (
|
||||||
req.bootstrap_room is not None
|
req.bootstrap_room is not None
|
||||||
), "req.bootstrap_room should not be None. Do not send requests directly to prefill or decode instances, but send to the router instead."
|
), "req.bootstrap_room should not be None. Do not send requests directly to prefill or decode instances, but send to the router instead."
|
||||||
|
|||||||
@@ -602,6 +602,8 @@ class ServerArgs:
|
|||||||
disaggregation_prefill_pp: Optional[int] = 1
|
disaggregation_prefill_pp: Optional[int] = 1
|
||||||
disaggregation_ib_device: Optional[str] = None
|
disaggregation_ib_device: Optional[str] = None
|
||||||
disaggregation_decode_enable_offload_kvcache: bool = False
|
disaggregation_decode_enable_offload_kvcache: bool = False
|
||||||
|
# Enable auto FAKE mode for decode node testing, no need to pass bootstrap_host in request
|
||||||
|
disaggregation_decode_enable_fake_auto: bool = False
|
||||||
num_reserved_decode_tokens: int = 512 # used for decode kv cache offload in PD
|
num_reserved_decode_tokens: int = 512 # used for decode kv cache offload in PD
|
||||||
# FIXME: hack to reduce ITL when decode bs is small
|
# FIXME: hack to reduce ITL when decode bs is small
|
||||||
disaggregation_decode_polling_interval: int = 1
|
disaggregation_decode_polling_interval: int = 1
|
||||||
@@ -4261,6 +4263,12 @@ class ServerArgs:
|
|||||||
action="store_true",
|
action="store_true",
|
||||||
help="Enable async KV cache offloading on decode server (PD mode).",
|
help="Enable async KV cache offloading on decode server (PD mode).",
|
||||||
)
|
)
|
||||||
|
parser.add_argument(
|
||||||
|
"--disaggregation-decode-enable-fake-auto",
|
||||||
|
action="store_true",
|
||||||
|
help="Auto enable FAKE mode for decode node testing, "
|
||||||
|
"no need to pass bootstrap_host and bootstrap_room in request.",
|
||||||
|
)
|
||||||
parser.add_argument(
|
parser.add_argument(
|
||||||
"--num-reserved-decode-tokens",
|
"--num-reserved-decode-tokens",
|
||||||
type=int,
|
type=int,
|
||||||
|
|||||||
Reference in New Issue
Block a user