| """Solo-only synthetic top32 qualification and graph timing; no serve requests.""" |
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import os |
| import resource |
| from collections.abc import Callable |
| from pathlib import Path |
|
|
| import torch |
|
|
| from ..sampled import MASS, Top32 |
|
|
| UUID = "GPU-bc62df2a-59cf-9828-7d43-558d1355bab7" |
|
|
|
|
| def timing(fn: Callable[[], object], repetitions: int = 100) -> dict[str, float]: |
| for _ in range(3): |
| fn() |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph): |
| for _ in range(10): |
| fn() |
| samples: list[float] = [] |
| start, stop = torch.cuda.Event(enable_timing=True), torch.cuda.Event(enable_timing=True) |
| for _ in range(repetitions): |
| start.record() |
| graph.replay() |
| stop.record() |
| stop.synchronize() |
| samples.append(start.elapsed_time(stop) / 10) |
| samples.sort() |
| return {"median_ms": samples[len(samples)//2], "p95_ms": samples[int(.95*len(samples))]} |
|
|
|
|
| def main() -> None: |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--out", type=Path, required=True) |
| a = p.parse_args() |
| if os.environ.get("CUDA_VISIBLE_DEVICES") != UUID: |
| raise RuntimeError("PP3 solo UUID required") |
| import fcntl |
| fd = int(os.environ["PP3_LOCK_FD"]) |
| fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB) |
| if str(torch.cuda.get_device_properties(0).uuid).removeprefix("GPU-") != UUID.removeprefix("GPU-"): |
| raise RuntimeError("wrong device") |
| torch.set_num_threads(8) |
| torch.set_num_interop_threads(1) |
| torch.manual_seed(4103) |
| dev = torch.device("cuda:0") |
| out: dict[str, object] = {"schema": "pp3.producer.v1", "classification": "MEASURED solo synthetic logits", |
| "complete": False, "torch": torch.__version__, "uuid": UUID, |
| "source_sha256": hashlib.sha256(Path(__file__).read_bytes()).hexdigest()} |
| out["producer_sha256"] = hashlib.sha256((Path(__file__).parents[1]/'sampled.py').read_bytes()).hexdigest() |
| rows: list[dict[str, object]] = [] |
| out["rows"] = rows |
|
|
| def save() -> None: |
| out["maxrss_kib"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss |
| out["gpu_peak_bytes"] = torch.cuda.max_memory_allocated() |
| a.out.write_text(json.dumps(out, indent=2)+"\n") |
|
|
| for batch in range(1, 5): |
| x = torch.randn(batch, 129280, device=dev).to(torch.bfloat16) |
| keys = torch.tensor([[4103+j, 1, 32768] for j in range(batch)], device=dev, dtype=torch.int64) |
| tau = torch.full((batch,), .7, device=dev) |
| sampler = Top32(batch, 129280, dev) |
| tok, ids, q = sampler(x, keys, tau) |
| torch.cuda.synchronize() |
| if sampler.error.any().item(): |
| raise AssertionError("valid proposal refused") |
| expected = torch.argsort(x.float(), descending=True, stable=True)[:, :32].sort(-1).values |
| torch.testing.assert_close(ids.long(), expected, rtol=0, atol=0) |
| if not torch.equal((q.double()*MASS).sum(-1), torch.full((batch,), MASS, device=dev, dtype=torch.float64)): |
| raise AssertionError("nonconserving mass") |
| if not ((ids == tok[:, None]) & (q > 0)).any(-1).all().item(): |
| raise AssertionError("proposal outside support") |
| for depth in (4, 5): |
| def sampled(depth: int = depth, sampler: Top32 = sampler, x: torch.Tensor = x, |
| keys: torch.Tensor = keys, tau: torch.Tensor = tau) -> None: |
| for _ in range(depth): |
| sampler(x, keys, tau) |
|
|
| def greedy(depth: int = depth, x: torch.Tensor = x) -> None: |
| for _ in range(depth): |
| x.argmax(-1) |
|
|
| row: dict[str, object] = {"batch": batch, "depth": depth, "sampled": timing(sampled), |
| "argmax": timing(greedy), "scope": "logit selection only; excludes Markov, TP exchange, wire"} |
| rows.append(row) |
| save() |
| print(json.dumps(row), flush=True) |
| out["complete"] = True |
| save() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
