[PD] Diversify fake-prefill handoff tokens (#37302)
Co-authored-by: xingyuliu <xingyuliu@fb.com> Co-authored-by: Lianmin Zheng <lianminzheng@gmail.com>
This commit is contained in:
co-authored by
xingyuliu
Lianmin Zheng
parent
db017e3490
commit
8b596c10b0
@@ -20,6 +20,7 @@ Life cycle of a request in the decode server
|
|||||||
|
|
||||||
from __future__ import annotations
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import hashlib
|
||||||
import logging
|
import logging
|
||||||
import time
|
import time
|
||||||
from collections import deque
|
from collections import deque
|
||||||
@@ -2012,6 +2013,22 @@ def alloc_for_decode_prealloc(
|
|||||||
return kv_loc
|
return kv_loc
|
||||||
|
|
||||||
|
|
||||||
|
def _generate_fake_prefill_handoff_output_id(req: Req) -> int:
|
||||||
|
"""Generate a deterministic synthetic handoff token for fake-PD requests.
|
||||||
|
|
||||||
|
Fake-PD skips prefill and KV transfer, so its allocated prompt KV has no
|
||||||
|
request-specific state and its metadata output token remains zero. Reusing
|
||||||
|
token 0 for every request can unrealistically correlate decode workloads.
|
||||||
|
Hash the request ID so every TP rank gets the same request-diverse token.
|
||||||
|
"""
|
||||||
|
if req.vocab_size is None or req.vocab_size <= 0:
|
||||||
|
raise ValueError("Fake-PD requests require a positive vocabulary size")
|
||||||
|
digest = hashlib.blake2b(
|
||||||
|
req.rid.encode(), digest_size=8, person=b"fake-pd"
|
||||||
|
).digest()
|
||||||
|
return int.from_bytes(digest, byteorder="little") % req.vocab_size
|
||||||
|
|
||||||
|
|
||||||
class DecodeTransferQueue(DecodeHiCacheTransferMixin):
|
class DecodeTransferQueue(DecodeHiCacheTransferMixin):
|
||||||
"""
|
"""
|
||||||
Store the requests that is polling kv
|
Store the requests that is polling kv
|
||||||
@@ -2142,6 +2159,10 @@ class DecodeTransferQueue(DecodeHiCacheTransferMixin):
|
|||||||
if replayed_boundary:
|
if replayed_boundary:
|
||||||
committed_output_id = decode_req.req.pd_rebootstrap_forced_output_id
|
committed_output_id = decode_req.req.pd_rebootstrap_forced_output_id
|
||||||
decode_req.req.pd_rebootstrap_forced_output_id = None
|
decode_req.req.pd_rebootstrap_forced_output_id = None
|
||||||
|
elif _is_fake_transfer(decode_req.req):
|
||||||
|
committed_output_id = _generate_fake_prefill_handoff_output_id(
|
||||||
|
decode_req.req
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
committed_output_id = output_id[0].item()
|
committed_output_id = output_id[0].item()
|
||||||
decode_req.req.output_ids.append(committed_output_id)
|
decode_req.req.output_ids.append(committed_output_id)
|
||||||
|
|||||||
Reference in New Issue
Block a user