82 lines
3.3 KiB
Python
82 lines
3.3 KiB
Python
"""P0 scheduler: inject a pre-fixed per-request verification horizon."""
|
|
|
|
from __future__ import annotations
|
|
|
|
from typing import Any
|
|
|
|
from vllm.v1.core.sched.scheduler import Scheduler
|
|
from vllm.v1.outputs import DraftTokenIds
|
|
|
|
from .common import digest, histogram, log_event, plan_summary, requested_k
|
|
|
|
|
|
class P0Scheduler(Scheduler):
|
|
"""A Scheduler that changes only post-drafter candidate visibility.
|
|
|
|
EAGLE has already produced its Kmax candidates when this method runs. The
|
|
P0 experiment therefore tests verifier-side ragged execution and collective
|
|
phase behavior, not a drafter speedup.
|
|
"""
|
|
|
|
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
|
super().__init__(*args, **kwargs)
|
|
self._p0_schedule_epoch = 0
|
|
self._p0_candidate_epoch = 0
|
|
|
|
def _p0_rank_fields(self) -> dict[str, Any]:
|
|
config = self.vllm_config.parallel_config
|
|
return {
|
|
"data_parallel_rank": getattr(config, "data_parallel_rank", None),
|
|
"data_parallel_size": getattr(config, "data_parallel_size", None),
|
|
"tensor_parallel_size": getattr(config, "tensor_parallel_size", None),
|
|
"expert_parallel_size": getattr(config, "expert_parallel_size", None),
|
|
}
|
|
|
|
def schedule(self, *args: Any, **kwargs: Any): # type: ignore[no-untyped-def]
|
|
output = super().schedule(*args, **kwargs)
|
|
self._p0_schedule_epoch += 1
|
|
log_event(
|
|
"scheduler",
|
|
"schedule",
|
|
core_local_epoch=self._p0_schedule_epoch,
|
|
**self._p0_rank_fields(),
|
|
**plan_summary(output),
|
|
)
|
|
return output
|
|
|
|
def update_draft_token_ids(self, draft_token_ids: DraftTokenIds) -> None:
|
|
self._p0_candidate_epoch += 1
|
|
before = [list(ids) for ids in draft_token_ids.draft_token_ids]
|
|
chosen: list[int] = []
|
|
after: list[list[int]] = []
|
|
structured_output_request_ids: list[str] = []
|
|
for request_id, ids in zip(draft_token_ids.req_ids, before):
|
|
request = self.requests.get(request_id)
|
|
if request is not None and getattr(request, "structured_output_request", None) is not None:
|
|
# Do not make grammar behavior part of this premise experiment.
|
|
chosen.append(len(ids))
|
|
after.append(ids)
|
|
structured_output_request_ids.append(str(request_id))
|
|
continue
|
|
k = requested_k(str(request_id), len(ids))
|
|
chosen.append(k)
|
|
after.append(ids[:k])
|
|
pairs = [
|
|
{"request_id": str(request_id), "before_k": len(ids), "after_k": after_k}
|
|
for request_id, ids, after_k in zip(draft_token_ids.req_ids, before, chosen)
|
|
]
|
|
log_event(
|
|
"scheduler",
|
|
"candidate_truncate",
|
|
candidate_epoch=self._p0_candidate_epoch,
|
|
**self._p0_rank_fields(),
|
|
candidate_count=len(pairs),
|
|
candidate_digest=digest(pairs),
|
|
before_k_histogram=histogram([item["before_k"] for item in pairs]),
|
|
after_k_histogram=histogram(chosen),
|
|
structured_output_request_ids=structured_output_request_ids,
|
|
)
|
|
super().update_draft_token_ids(
|
|
DraftTokenIds(req_ids=draft_token_ids.req_ids, draft_token_ids=after)
|
|
)
|