feat: add agentic pd hybrid benchmark prototype
This commit is contained in:
235
src/agentic_pd_hybrid/policies.py
Normal file
235
src/agentic_pd_hybrid/policies.py
Normal file
@@ -0,0 +1,235 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import Counter
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Protocol
|
||||
|
||||
from agentic_pd_hybrid.topology import SingleNodeTopology
|
||||
from agentic_pd_hybrid.trace import TraceRequest
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionRouteState:
|
||||
last_decode_worker: str | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class RoutingDecision:
|
||||
policy_name: str
|
||||
prefill_worker_id: str
|
||||
decode_worker_id: str
|
||||
decode_worker_index: int
|
||||
reuse_expected: bool
|
||||
observed_overlap_blocks: int
|
||||
kv_transfer_blocks: int
|
||||
inflight_decode_load: int
|
||||
session_id: str
|
||||
request_id: str
|
||||
turn_id: int
|
||||
|
||||
@property
|
||||
def observed_reuse(self) -> bool:
|
||||
return self.observed_overlap_blocks > 0
|
||||
|
||||
@property
|
||||
def re_prefill_required(self) -> bool:
|
||||
return self.turn_id > 1 and self.observed_overlap_blocks == 0
|
||||
|
||||
|
||||
@dataclass
|
||||
class RoutingState:
|
||||
prefill_cursor: int = 0
|
||||
decode_cursor: int = 0
|
||||
session_state: dict[str, SessionRouteState] = field(default_factory=dict)
|
||||
inflight_decode: Counter[str] = field(default_factory=Counter)
|
||||
decode_assignment_counts: Counter[str] = field(default_factory=Counter)
|
||||
decode_resident_blocks: dict[str, set[int]] = field(default_factory=dict)
|
||||
|
||||
@classmethod
|
||||
def create(cls, topology: SingleNodeTopology) -> "RoutingState":
|
||||
return cls(
|
||||
decode_resident_blocks={
|
||||
worker.worker_id: set() for worker in topology.route_workers
|
||||
}
|
||||
)
|
||||
|
||||
def next_prefill_worker_id(self, topology: SingleNodeTopology) -> str:
|
||||
if not topology.prefill_workers:
|
||||
return "none"
|
||||
worker = topology.prefill_workers[self.prefill_cursor % len(topology.prefill_workers)]
|
||||
self.prefill_cursor += 1
|
||||
return worker.worker_id
|
||||
|
||||
def next_decode_worker_id(self, topology: SingleNodeTopology) -> str:
|
||||
route_workers = topology.route_workers
|
||||
worker = route_workers[self.decode_cursor % len(route_workers)]
|
||||
self.decode_cursor += 1
|
||||
return worker.worker_id
|
||||
|
||||
def finish(self, request: TraceRequest, decision: RoutingDecision) -> None:
|
||||
session = self.session_state.setdefault(request.session_id, SessionRouteState())
|
||||
session.last_decode_worker = decision.decode_worker_id
|
||||
self.decode_resident_blocks[decision.decode_worker_id].update(request.hash_ids)
|
||||
self.inflight_decode[decision.decode_worker_id] -= 1
|
||||
if self.inflight_decode[decision.decode_worker_id] <= 0:
|
||||
del self.inflight_decode[decision.decode_worker_id]
|
||||
|
||||
|
||||
class RoutingPolicy(Protocol):
|
||||
name: str
|
||||
|
||||
def select(
|
||||
self,
|
||||
request: TraceRequest,
|
||||
*,
|
||||
topology: SingleNodeTopology,
|
||||
state: RoutingState,
|
||||
) -> RoutingDecision:
|
||||
...
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class DefaultPolicy:
|
||||
name: str = "default"
|
||||
|
||||
def select(
|
||||
self,
|
||||
request: TraceRequest,
|
||||
*,
|
||||
topology: SingleNodeTopology,
|
||||
state: RoutingState,
|
||||
) -> RoutingDecision:
|
||||
prefill_worker_id = state.next_prefill_worker_id(topology)
|
||||
decode_worker_id = state.next_decode_worker_id(topology)
|
||||
return _build_decision(
|
||||
policy_name=self.name,
|
||||
request=request,
|
||||
topology=topology,
|
||||
state=state,
|
||||
prefill_worker_id=prefill_worker_id,
|
||||
decode_worker_id=decode_worker_id,
|
||||
reuse_expected=False,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StickyDecodePolicy:
|
||||
name: str = "sticky"
|
||||
|
||||
def select(
|
||||
self,
|
||||
request: TraceRequest,
|
||||
*,
|
||||
topology: SingleNodeTopology,
|
||||
state: RoutingState,
|
||||
) -> RoutingDecision:
|
||||
session = state.session_state.get(request.session_id)
|
||||
prefill_worker_id = state.next_prefill_worker_id(topology)
|
||||
if request.turn_id > 1 and session and session.last_decode_worker is not None:
|
||||
decode_worker_id = session.last_decode_worker
|
||||
reuse_expected = True
|
||||
else:
|
||||
decode_worker_id = state.next_decode_worker_id(topology)
|
||||
reuse_expected = False
|
||||
return _build_decision(
|
||||
policy_name=self.name,
|
||||
request=request,
|
||||
topology=topology,
|
||||
state=state,
|
||||
prefill_worker_id=prefill_worker_id,
|
||||
decode_worker_id=decode_worker_id,
|
||||
reuse_expected=reuse_expected,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class KvAwarePolicy:
|
||||
name: str = "kv-aware"
|
||||
sticky_bonus: int = 1
|
||||
|
||||
def select(
|
||||
self,
|
||||
request: TraceRequest,
|
||||
*,
|
||||
topology: SingleNodeTopology,
|
||||
state: RoutingState,
|
||||
) -> RoutingDecision:
|
||||
prefill_worker_id = state.next_prefill_worker_id(topology)
|
||||
session = state.session_state.get(request.session_id)
|
||||
|
||||
best_decode_worker_id: str | None = None
|
||||
best_score: tuple[int, int, int] | None = None
|
||||
for worker in topology.route_workers:
|
||||
overlap = _overlap_blocks(request, state, worker.worker_id)
|
||||
sticky = int(session is not None and session.last_decode_worker == worker.worker_id)
|
||||
inflight_penalty = -state.inflight_decode.get(worker.worker_id, 0)
|
||||
assignment_penalty = -state.decode_assignment_counts.get(worker.worker_id, 0)
|
||||
score = (
|
||||
overlap + sticky * self.sticky_bonus,
|
||||
sticky,
|
||||
inflight_penalty,
|
||||
assignment_penalty,
|
||||
)
|
||||
if best_score is None or score > best_score:
|
||||
best_score = score
|
||||
best_decode_worker_id = worker.worker_id
|
||||
|
||||
assert best_decode_worker_id is not None
|
||||
reuse_expected = bool(best_score and best_score[0] > 0)
|
||||
return _build_decision(
|
||||
policy_name=self.name,
|
||||
request=request,
|
||||
topology=topology,
|
||||
state=state,
|
||||
prefill_worker_id=prefill_worker_id,
|
||||
decode_worker_id=best_decode_worker_id,
|
||||
reuse_expected=reuse_expected,
|
||||
)
|
||||
|
||||
|
||||
def create_policy(name: str) -> RoutingPolicy:
|
||||
normalized = name.strip().lower()
|
||||
if normalized == "default":
|
||||
return DefaultPolicy()
|
||||
if normalized == "sticky":
|
||||
return StickyDecodePolicy()
|
||||
if normalized in {"kv-aware", "kv_aware", "kv"}:
|
||||
return KvAwarePolicy()
|
||||
raise ValueError(f"Unsupported policy: {name}")
|
||||
|
||||
|
||||
def _build_decision(
|
||||
*,
|
||||
policy_name: str,
|
||||
request: TraceRequest,
|
||||
topology: SingleNodeTopology,
|
||||
state: RoutingState,
|
||||
prefill_worker_id: str,
|
||||
decode_worker_id: str,
|
||||
reuse_expected: bool,
|
||||
) -> RoutingDecision:
|
||||
overlap = _overlap_blocks(request, state, decode_worker_id)
|
||||
state.inflight_decode[decode_worker_id] += 1
|
||||
state.decode_assignment_counts[decode_worker_id] += 1
|
||||
return RoutingDecision(
|
||||
policy_name=policy_name,
|
||||
prefill_worker_id=prefill_worker_id,
|
||||
decode_worker_id=decode_worker_id,
|
||||
decode_worker_index=topology.route_index(decode_worker_id),
|
||||
reuse_expected=reuse_expected,
|
||||
observed_overlap_blocks=overlap,
|
||||
kv_transfer_blocks=max(0, len(request.hash_ids) - overlap),
|
||||
inflight_decode_load=state.inflight_decode[decode_worker_id],
|
||||
session_id=request.session_id,
|
||||
request_id=request.request_id,
|
||||
turn_id=request.turn_id,
|
||||
)
|
||||
|
||||
|
||||
def _overlap_blocks(
|
||||
request: TraceRequest,
|
||||
state: RoutingState,
|
||||
decode_worker_id: str,
|
||||
) -> int:
|
||||
resident = state.decode_resident_blocks.get(decode_worker_id, set())
|
||||
return sum(1 for block in request.hash_ids if block in resident)
|
||||
Reference in New Issue
Block a user