[PD] Tiny register info field cleanup for mooncake backend (#22016)
This commit is contained in:
@@ -81,6 +81,7 @@ class TransferInfo:
|
||||
dst_state_indices: List[int]
|
||||
required_dst_info_num: int
|
||||
is_dummy: bool
|
||||
# Note: always put the optional staging field at the final (it will be set through 'STAGING_RSP' pkg when needed)
|
||||
staging: Optional[StagingTransferInfo] = None
|
||||
|
||||
@classmethod
|
||||
@@ -129,6 +130,7 @@ class KVArgsRegisterInfo:
|
||||
dst_state_dim_per_tensor: list[int]
|
||||
# HiSparse: decode host pool stores KV at token granularity
|
||||
enable_hisparse: bool = False
|
||||
# Note: always put the staging field at the final (since the staging field is optional and contains multiple inputs)
|
||||
staging: Optional[StagingRegisterInfo] = None
|
||||
|
||||
@classmethod
|
||||
@@ -157,6 +159,7 @@ class KVArgsRegisterInfo:
|
||||
enable_hisparse=(
|
||||
msg[12].decode("ascii") == "1" if len(msg) > 12 else False
|
||||
),
|
||||
# Note: always put the staging field at the final
|
||||
staging=StagingRegisterInfo.from_zmq_fields(msg, 13),
|
||||
)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user