| """Solo synthetic qualification of fused Markov/q/graph semantics (no weights).""" |
| from __future__ import annotations |
|
|
| import argparse |
| import fcntl |
| import hashlib |
| import json |
| import os |
| import resource |
| import time |
| from collections.abc import Callable |
| from pathlib import Path |
| from types import SimpleNamespace |
| from typing import Any, cast |
| from unittest.mock import patch |
|
|
| import torch |
|
|
| from ...worker import sampled_fused |
| from ...worker.sampled import MASS, Top32 |
| from ...worker.sampled_fused import FusedContext |
| from ...worker.tp2 import MarkovShard |
|
|
| UUID = "GPU-bc62df2a-59cf-9828-7d43-558d1355bab7" |
|
|
|
|
| def digest(path: Path) -> str: |
| return hashlib.sha256(path.read_bytes()).hexdigest() |
|
|
|
|
| def model_fixture() -> Any: |
| dev = torch.device("cuda:0") |
| vocab, rank_dim = 129280, 256 |
| w1 = (torch.randn(vocab, rank_dim, device=dev) / 8).bfloat16() |
| slots = torch.full((vocab,), -1, dtype=torch.int64, device=dev) |
| slots[:4] = torch.arange(4, device=dev) |
| ranks: list[Any] = [] |
| for rank, (lo, hi) in enumerate(((0, 64768), (64768, vocab))): |
| mb = MarkovShard.__new__(MarkovShard) |
| mb.w1 = w1 |
| mb.w2 = (torch.randn(hi - lo, rank_dim, device=dev) / 8).bfloat16() |
| mb.slot_of_token = slots |
| mb.cache = (torch.randn(4, hi - lo, device=dev) / 4).bfloat16() |
| ranks.append(SimpleNamespace(r=rank, v0=lo, v1=hi, markov=mb, dev=dev)) |
| return SimpleNamespace(ALL=(0, 1), devs=[dev, dev], ranks=ranks, emulated=True) |
|
|
|
|
| def qualify(model: Any, batch: int, depth: int, mask: torch.Tensor) -> dict[str, Any]: |
| dev = model.devs[0] |
| ctx = FusedContext(model, batch, depth) |
| keys = torch.tensor([[4104 + j, 1, 32768] for j in range(batch)], device=dev, dtype=torch.int64) |
| ctx.update(keys, mask) |
| logits = {r: torch.randn(batch, depth, rk.v1 - rk.v0, device=dev).bfloat16() |
| for r, rk in enumerate(model.ranks)} |
| output = {r: torch.zeros(batch, depth, device=dev, dtype=torch.int64) for r in model.ALL} |
| sampler = Top32(batch, 129280, dev) |
| checks = 0 |
| |
| |
| for initial in (0, 7): |
| anchor = torch.tensor([initial if j % 2 == 0 else 2 for j in range(batch)], device=dev) |
| for pos in range(depth): |
| previous = {r: anchor if pos == 0 else output[r][:, pos - 1] for r in model.ALL} |
| ctx.select(model.ALL, logits, previous, output, pos) |
| support, q = ctx.result() |
| parts = [] |
| for r in model.ALL: |
| mb = model.ranks[r].markov |
| prev = previous[r].contiguous() |
| parts.append((logits[r][:, pos] + mb.bias(prev, mb.embed(prev))).bfloat16()) |
| full = torch.cat(parts, -1) |
| poskeys = keys.clone() |
| poskeys[:, 2] += pos |
| token, ref_ids, ref_q = sampler(full, poskeys, torch.full((batch,), .7, device=dev)) |
| expected = torch.where(mask, token.long(), full.argmax(-1)) |
| torch.testing.assert_close(output[0][:, pos], expected, rtol=0, atol=0) |
| torch.testing.assert_close(output[1][:, pos], expected, rtol=0, atol=0) |
| torch.testing.assert_close(support[:, pos], ref_ids, rtol=0, atol=0) |
| torch.testing.assert_close(q[:, pos], ref_q * mask[:, None], rtol=0, atol=0) |
| if not torch.equal((q[:, pos].double() * MASS).sum(-1), mask.double() * MASS): |
| raise AssertionError("dyadic mass differs") |
| checks += batch |
| |
| |
| anchor = torch.zeros(batch, device=dev, dtype=torch.int64) |
|
|
| def run() -> None: |
| for pos in range(depth): |
| previous = {r: anchor if pos == 0 else output[r][:, pos - 1] for r in model.ALL} |
| ctx.select(model.ALL, logits, previous, output, pos) |
|
|
| for _ in range(2): |
| run() |
| torch.cuda.synchronize() |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph): |
| run() |
| times: list[float] = [] |
| event_factory = cast(Callable[..., torch.cuda.Event], torch.cuda.Event) |
| start, end = event_factory(enable_timing=True), event_factory(enable_timing=True) |
| for rnd in range(100): |
| ctx.state[0]["keys"][:, 1].fill_(rnd) |
| start.record() |
| graph.replay() |
| end.record() |
| end.synchronize() |
| ctx.result() |
| times.append(start.elapsed_time(end)) |
| |
| before = (output[0].clone(), ctx.state[0]["ids"].clone(), ctx.state[0]["q"].clone()) |
| run() |
| for actual, expected in zip((output[0], ctx.state[0]["ids"], ctx.state[0]["q"]), before, strict=True): |
| torch.testing.assert_close(actual, expected, rtol=0, atol=0) |
| times.sort() |
| return {"batch": batch, "depth": depth, "sampled_rows": int(mask.sum().item()), |
| "exact_position_checks": checks, "graph_replays": 100, |
| "selector_median_ms": times[50], "selector_p95_ms": times[95], |
| "scope": "MEASURED two logical TP2 shards on solo; selector including Markov, excludes backbone/pair"} |
|
|
|
|
| def compile_peer_branches(model: Any) -> None: |
| """Compile both real-rank specializations without executing peer traffic. |
| |
| This catches JIT/assembly errors; it does not qualify physical-pair numerics, |
| visibility, timeout or graph interoperability. |
| """ |
| |
| kernel: Any = sampled_fused._kernel |
|
|
| class CompileOnly: |
| def warmup(self, *args: Any, **kwargs: Any) -> Any: |
| return kernel.warmup(*args, **kwargs) |
|
|
| def __getitem__(self, grid: Any) -> Any: |
| def compile_call(*args: Any, **kwargs: Any) -> Any: |
| return kernel.warmup(*args, grid=grid, **kwargs) |
| return compile_call |
|
|
| model.emulated = False |
| try: |
| ctx = FusedContext(model, 1, 5) |
| logits = {r: torch.zeros(1, 5, rk.v1 - rk.v0, device=rk.dev, dtype=torch.bfloat16) |
| for r, rk in enumerate(model.ranks)} |
| previous = {r: torch.zeros(1, device=rk.dev, dtype=torch.int64) for r, rk in enumerate(model.ranks)} |
| outputs = {r: torch.zeros(1, 5, device=rk.dev, dtype=torch.int64) for r, rk in enumerate(model.ranks)} |
| with patch.object(sampled_fused, "_kernel", CompileOnly()): |
| ctx.select((0, 1), logits, previous, outputs, 0) |
| finally: |
| model.emulated = True |
|
|
|
|
| def edge_cases(model: Any) -> dict[str, bool]: |
| dev = model.devs[0] |
| ctx = FusedContext(model, 2, 5) |
| keys = torch.tensor([[4104, 7, 100], [4105, 8, 101]], device=dev, dtype=torch.int64) |
| mask = torch.tensor([True, False], device=dev) |
| ctx.update(keys, mask) |
| logits = {r: torch.zeros(2, 5, rk.v1 - rk.v0, device=dev, dtype=torch.bfloat16) |
| for r, rk in enumerate(model.ranks)} |
| previous = {r: torch.zeros(2, device=dev, dtype=torch.int64) for r in model.ALL} |
| output = {r: torch.zeros(2, 5, device=dev, dtype=torch.int64) for r in model.ALL} |
| |
| |
| saved = [rk.markov.cache[0].clone() for rk in model.ranks] |
| try: |
| for rk in model.ranks: |
| rk.markov.cache[0].zero_() |
| logits[0][:, 0, 1] = -0. |
| ctx.select(model.ALL, logits, previous, output, 0) |
| ids, q = ctx.result() |
| expected_ids = torch.arange(32, device=dev, dtype=torch.int32).expand(2, 32) |
| torch.testing.assert_close(ids[:, 0], expected_ids, rtol=0, atol=0) |
| assert output[0][1, 0].item() == 0 |
| torch.testing.assert_close(q[0, 0], torch.full((32,), 1 / 32, device=dev), rtol=0, atol=0) |
| assert torch.count_nonzero(q[1, 0]).item() == 0 |
| for bad in (float("nan"), float("inf")): |
| logits[1][0, 0, 3] = bad |
| ctx.select(model.ALL, logits, previous, output, 0) |
| for pos in range(1, 5): |
| continuation = {r: output[r][:, pos - 1] for r in model.ALL} |
| ctx.select(model.ALL, logits, continuation, output, pos) |
| try: |
| ctx.result() |
| except RuntimeError: |
| pass |
| else: |
| raise AssertionError("invalid fused logit accepted") |
| for rank in model.ALL: |
| logits[rank][:, 0].fill_(float("-inf")) |
| ctx.select(model.ALL, logits, previous, output, 0) |
| try: |
| ctx.result() |
| except RuntimeError: |
| pass |
| else: |
| raise AssertionError("all-masked fused logits accepted") |
| finally: |
| for rk, cache in zip(model.ranks, saved, strict=True): |
| rk.markov.cache[0].copy_(cache) |
| return {"lower_id_ties": True, "signed_zero": True, "mixed_t0_no_draw": True, |
| "invalid_intermediate_safe_continuation": True, |
| "nan_inf_all_masked_refused": True} |
|
|
|
|
| def main() -> None: |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--out", type=Path, required=True) |
| args = parser.parse_args() |
| if os.environ.get("CUDA_VISIBLE_DEVICES") != UUID: |
| raise RuntimeError("PP4 assigned solo UUID required") |
| fd = int(os.environ["PP4_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 physical GPU") |
| torch.set_num_threads(8) |
| torch.set_num_interop_threads(1) |
| torch.manual_seed(4104) |
| rows: list[dict[str, Any]] = [] |
| root = Path(__file__).resolve().parents[2] |
| result: dict[str, Any] = {"schema": "pp4.fused-qualification.v1", "complete": False, |
| "classification": "MEASURED synthetic solo", "uuid": UUID, |
| "torch": torch.__version__, "rows": rows, |
| "source_sha256": digest(Path(__file__)), |
| "fused_sha256": digest(root / "worker/sampled_fused.py")} |
| started = time.monotonic() |
| model = model_fixture() |
| compile_peer_branches(model) |
| result["peer_specializations_compile_only"] = True |
| result["edges"] = edge_cases(model) |
| for batch in range(1, 5): |
| for depth in (4, 5): |
| for mixed in (False, True): |
| mask = torch.tensor([j % 2 == 0 if mixed else True for j in range(batch)], device=model.devs[0]) |
| row = qualify(model, batch, depth, mask) |
| rows.append(row) |
| result["maxrss_kib"] = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss |
| result["gpu_peak_bytes"] = torch.cuda.max_memory_allocated() |
| args.out.write_text(json.dumps(result, indent=2) + "\n") |
| print(json.dumps(row), flush=True) |
| result["elapsed_s"] = time.monotonic() - started |
| result["complete"] = True |
| args.out.write_text(json.dumps(result, indent=2) + "\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
