[PD] support pd fake transfer for warmup (#5726)

This commit is contained in:
ybyang
2025-04-29 00:33:20 +08:00
committed by GitHub
parent 92ab0a2055
commit c6c6264073
6 changed files with 146 additions and 7 deletions

View File

@@ -29,6 +29,7 @@ import torch
from sglang.srt.disaggregation.base import BaseKVManager, KVArgs, KVPoll
from sglang.srt.disaggregation.utils import (
DisaggregationMode,
FakeBootstrapHost,
KVClassType,
ReqToMetadataIdxAllocator,
TransferBackend,
@@ -116,7 +117,11 @@ class PrefillBootstrapQueue:
return kv_manager
def add(self, req: Req) -> None:
kv_sender_class = get_kv_class(self.transfer_backend, KVClassType.SENDER)
if req.bootstrap_host == FakeBootstrapHost:
# Fake transfer for warmup reqs
kv_sender_class = get_kv_class(TransferBackend.FAKE, KVClassType.SENDER)
else:
kv_sender_class = get_kv_class(self.transfer_backend, KVClassType.SENDER)
req.disagg_kv_sender = kv_sender_class(
mgr=self.kv_manager,
bootstrap_addr=f"{req.bootstrap_host}:{self.bootstrap_port}",