Add CollectiveSpec P0 phase trace harness

scripts/collectivespec/p0_overlay/collectivespec_p0/scheduler.py

@@ -0,0 +1,81 @@
+"""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)
+        )