[PD] support pd fake transfer for warmup (#5726)
This commit is contained in:
@@ -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}",
|
||||
|
||||
Reference in New Issue
Block a user