Add CollectiveSpec P0 phase trace harness
This commit is contained in:
@@ -0,0 +1,6 @@
|
||||
"""Non-invasive instrumentation for the CollectiveSpec P0 premise test.
|
||||
|
||||
This package is loaded only when its parent ``p0_overlay`` directory is added
|
||||
to ``PYTHONPATH`` and vLLM is explicitly given the custom scheduler/worker
|
||||
classes. It is intentionally not imported by normal aituner runs.
|
||||
"""
|
||||
158
scripts/collectivespec/p0_overlay/collectivespec_p0/common.py
Normal file
158
scripts/collectivespec/p0_overlay/collectivespec_p0/common.py
Normal file
@@ -0,0 +1,158 @@
|
||||
"""Small, dependency-free helpers shared by the P0 vLLM extensions."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import os
|
||||
import threading
|
||||
import time
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
_LOCK = threading.Lock()
|
||||
_POLICY: dict[str, Any] | None = None
|
||||
|
||||
|
||||
def _canonical(value: Any) -> str:
|
||||
return json.dumps(value, ensure_ascii=True, separators=(",", ":"), sort_keys=True)
|
||||
|
||||
|
||||
def digest(value: Any) -> str:
|
||||
return hashlib.blake2b(_canonical(value).encode("utf-8"), digest_size=16).hexdigest()
|
||||
|
||||
|
||||
def policy() -> dict[str, Any]:
|
||||
global _POLICY
|
||||
if _POLICY is not None:
|
||||
return _POLICY
|
||||
path_value = os.environ.get("COLLECTIVESPEC_P0_POLICY_PATH")
|
||||
if not path_value:
|
||||
raise RuntimeError("COLLECTIVESPEC_P0_POLICY_PATH is required for P0")
|
||||
path = Path(path_value)
|
||||
value = json.loads(path.read_text(encoding="utf-8"))
|
||||
if not isinstance(value, dict):
|
||||
raise RuntimeError("P0 policy must be a JSON object")
|
||||
max_k = value.get("max_k")
|
||||
seed = value.get("seed")
|
||||
kind = value.get("policy")
|
||||
if isinstance(max_k, bool) or not isinstance(max_k, int) or max_k < 0:
|
||||
raise RuntimeError("P0 policy max_k must be a non-negative integer")
|
||||
if isinstance(seed, bool) or not isinstance(seed, int):
|
||||
raise RuntimeError("P0 policy seed must be an integer")
|
||||
if kind not in {"blake2b-request-static-k", "constant-k"}:
|
||||
raise RuntimeError(f"Unsupported P0 policy: {kind!r}")
|
||||
if kind == "constant-k":
|
||||
constant_k = value.get("constant_k")
|
||||
if isinstance(constant_k, bool) or not isinstance(constant_k, int):
|
||||
raise RuntimeError("constant-k policy requires integer constant_k")
|
||||
if not 0 <= constant_k <= max_k:
|
||||
raise RuntimeError("constant_k must be within [0, max_k]")
|
||||
_POLICY = value
|
||||
return value
|
||||
|
||||
|
||||
def requested_k(request_id: str, candidate_count: int) -> int:
|
||||
value = policy()
|
||||
max_k = int(value["max_k"])
|
||||
if candidate_count < 0:
|
||||
raise RuntimeError("candidate_count must be non-negative")
|
||||
if value["policy"] == "constant-k":
|
||||
chosen = int(value["constant_k"])
|
||||
else:
|
||||
material = f"{value['seed']}|{request_id}".encode("utf-8")
|
||||
chosen = int.from_bytes(hashlib.blake2b(material, digest_size=8).digest(), "big") % (max_k + 1)
|
||||
return min(chosen, candidate_count)
|
||||
|
||||
|
||||
def histogram(values: list[int]) -> dict[str, int]:
|
||||
counts: dict[str, int] = {}
|
||||
for value in values:
|
||||
key = str(int(value))
|
||||
counts[key] = counts.get(key, 0) + 1
|
||||
return dict(sorted(counts.items(), key=lambda item: int(item[0])))
|
||||
|
||||
|
||||
def _jsonable(value: Any) -> Any:
|
||||
if value is None or isinstance(value, (str, int, float, bool)):
|
||||
return value
|
||||
if isinstance(value, (list, tuple)):
|
||||
return [_jsonable(item) for item in value]
|
||||
if isinstance(value, dict):
|
||||
return {str(key): _jsonable(item) for key, item in value.items()}
|
||||
tolist = getattr(value, "tolist", None)
|
||||
if callable(tolist):
|
||||
return _jsonable(tolist())
|
||||
item = getattr(value, "item", None)
|
||||
if callable(item):
|
||||
try:
|
||||
return _jsonable(item())
|
||||
except (TypeError, ValueError):
|
||||
pass
|
||||
return repr(value)
|
||||
|
||||
|
||||
def dp_metadata_summary(metadata: Any) -> dict[str, Any] | None:
|
||||
if metadata is None:
|
||||
return None
|
||||
keys = (
|
||||
"num_pad_tokens",
|
||||
"num_tokens_across_dp",
|
||||
"num_tokens_per_rank",
|
||||
"num_reqs_per_rank",
|
||||
"is_prompt_batch",
|
||||
"all_prefill",
|
||||
"all_decode",
|
||||
"should_ubatch",
|
||||
"input_fits_in_drafter",
|
||||
"cudagraph_mode",
|
||||
)
|
||||
return {key: _jsonable(getattr(metadata, key, None)) for key in keys}
|
||||
|
||||
|
||||
def plan_summary(scheduler_output: Any) -> dict[str, Any]:
|
||||
scheduled = getattr(scheduler_output, "num_scheduled_tokens", {}) or {}
|
||||
spec = getattr(scheduler_output, "scheduled_spec_decode_tokens", {}) or {}
|
||||
ordered = [
|
||||
{
|
||||
"request_id": str(request_id),
|
||||
"scheduled_tokens": int(tokens),
|
||||
"spec_k": len(spec.get(request_id, [])),
|
||||
}
|
||||
for request_id, tokens in scheduled.items()
|
||||
]
|
||||
semantic = sorted(ordered, key=lambda item: item["request_id"])
|
||||
spec_ks = [item["spec_k"] for item in ordered if item["spec_k"] > 0]
|
||||
return {
|
||||
"ordered_plan_digest": digest(ordered),
|
||||
"semantic_plan_digest": digest(semantic),
|
||||
"request_count": len(ordered),
|
||||
"total_scheduled_rows": int(getattr(scheduler_output, "total_num_scheduled_tokens", 0)),
|
||||
"spec_request_count": len(spec),
|
||||
"spec_k_histogram": histogram(spec_ks),
|
||||
"verify_logical_rows": sum(1 + item["spec_k"] for item in ordered if item["spec_k"] > 0),
|
||||
"scheduled_dp_metadata": dp_metadata_summary(
|
||||
getattr(scheduler_output, "scheduled_dp_metadata", None)
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
def log_event(role: str, event: str, **payload: Any) -> None:
|
||||
root_value = os.environ.get("COLLECTIVESPEC_P0_LOG_DIR")
|
||||
if not root_value:
|
||||
raise RuntimeError("COLLECTIVESPEC_P0_LOG_DIR is required for P0")
|
||||
root = Path(root_value)
|
||||
root.mkdir(parents=True, exist_ok=True)
|
||||
record = {
|
||||
"schema_version": "collectivespec-p0-log-v1",
|
||||
"event": event,
|
||||
"monotonic_s": time.monotonic(),
|
||||
"pid": os.getpid(),
|
||||
"policy_version": policy().get("version"),
|
||||
"role": role,
|
||||
**{key: _jsonable(value) for key, value in payload.items()},
|
||||
}
|
||||
path = root / f"{role}-pid{os.getpid()}.jsonl"
|
||||
with _LOCK, path.open("a", encoding="utf-8") as handle:
|
||||
handle.write(json.dumps(record, ensure_ascii=True, sort_keys=True) + "\n")
|
||||
@@ -0,0 +1,6 @@
|
||||
{
|
||||
"max_k": 3,
|
||||
"policy": "blake2b-request-static-k",
|
||||
"seed": 20260713,
|
||||
"version": "collectivespec-p0-v1"
|
||||
}
|
||||
@@ -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)
|
||||
)
|
||||
@@ -0,0 +1,44 @@
|
||||
"""P0 worker: record actual rank-side target execution phases."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
from vllm.v1.worker.gpu_worker import Worker
|
||||
|
||||
from .common import log_event, plan_summary
|
||||
|
||||
|
||||
class P0Worker(Worker):
|
||||
"""Log rank-local execution without modifying model or collective calls."""
|
||||
|
||||
def __init__(self, *args: Any, **kwargs: Any) -> None:
|
||||
super().__init__(*args, **kwargs)
|
||||
self._p0_execute_epoch = 0
|
||||
|
||||
def _p0_rank_fields(self) -> dict[str, Any]:
|
||||
config = self.vllm_config.parallel_config
|
||||
return {
|
||||
"global_rank": getattr(self, "rank", None),
|
||||
"local_rank": getattr(self, "local_rank", None),
|
||||
"data_parallel_rank": getattr(config, "data_parallel_rank", None),
|
||||
"tensor_parallel_rank": getattr(config, "tensor_parallel_rank", None),
|
||||
"expert_parallel_size": getattr(config, "expert_parallel_size", None),
|
||||
"tensor_parallel_size": getattr(config, "tensor_parallel_size", None),
|
||||
"data_parallel_size": getattr(config, "data_parallel_size", None),
|
||||
"speculative_kmax": getattr(
|
||||
getattr(self, "speculative_config", None), "num_speculative_tokens", None
|
||||
),
|
||||
}
|
||||
|
||||
def execute_model(self, scheduler_output: Any): # type: ignore[no-untyped-def]
|
||||
self._p0_execute_epoch += 1
|
||||
fields = {
|
||||
"worker_phase_epoch": self._p0_execute_epoch,
|
||||
**self._p0_rank_fields(),
|
||||
**plan_summary(scheduler_output),
|
||||
}
|
||||
log_event("worker", "target_execute_begin", **fields)
|
||||
output = super().execute_model(scheduler_output)
|
||||
log_event("worker", "target_execute_end", **fields)
|
||||
return output
|
||||
Reference in New Issue
Block a user