| |
| """Validate the local HF package and caller-owned deployment contract on CUDA.""" |
| from __future__ import annotations |
|
|
| import base64 |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import torch |
| from kernels import get_local_kernel |
|
|
| import benchmark as bench |
|
|
|
|
| def main() -> None: |
| here = Path(__file__).resolve().parent |
| runtime = here / "build/torch-cuda" |
| metadata = json.loads((runtime / "metadata.json").read_text()) |
| for filename, expected in metadata["digest"]["files"].items(): |
| actual = base64.b64encode(hashlib.sha256((runtime / filename).read_bytes()).digest()).decode() |
| assert actual == expected, f"digest mismatch: {filename}" |
| op = get_local_kernel(here, trust_remote_code=True) |
| assert all(hasattr(op, name) for name in ("Pool", "flush_inplace", "run")) |
| namespace = {} |
| exec(json.loads((here / "definition.json").read_text())["reference"], namespace) |
| reference = namespace["run"] |
| device = torch.device("cuda") |
| ns = 8 |
| pool = op.Pool(ns, device=device) |
| is_kda = hasattr(pool, "gbuf") |
| idx = torch.tensor([0, 7, 2, 5, 1, 3], dtype=torch.int32, device=device) |
| counters = torch.where(torch.arange(ns * 32, device=device).reshape(ns, 32) % 3 == 0, 16, 7).int() |
| selected = torch.zeros(ns, dtype=torch.bool, device=device) |
| selected[idx.long()] = True |
| selected[0] = False |
| due = selected[:, None] & (counters == 16) |
| fields = ["codes", "s_k", "s_v", "uq", "kbuf", "ubuf", "w", "hcnt"] + (["gbuf", "pcum"] if is_kda else ["p"]) |
| q0 = torch.empty(128, 4, device=device) |
|
|
| def restore(inputs): |
| pool.codes.copy_(inputs[0]); pool.s_k.copy_(inputs[1]); pool.s_v.copy_(inputs[2]) |
| pool.uq[..., :128].copy_(inputs[3]); pool.uq[..., 128:].copy_(inputs[4]) |
| pool.kbuf.copy_(inputs[5]); pool.ubuf.copy_(inputs[6]) |
| if is_kda: |
| pool.gbuf.copy_(inputs[7]); pool.w.copy_(inputs[8]); pool.pcum.copy_(inputs[9]) |
| else: |
| pool.w.copy_(inputs[7]); pool.p[..., 0].copy_(inputs[8]) |
| pool.hcnt.copy_(counters) |
| q0.copy_(inputs[-1]) |
|
|
| def snapshot(): |
| return {name: getattr(pool, name).clone() for name in fields} |
|
|
| def verify(inputs, before): |
| ref = reference(*[tensor.clone() for tensor in inputs]) |
| actual = (pool.codes, pool.s_k, pool.s_v, pool.uq[..., :128], pool.uq[..., 128:]) |
| |
| out = tuple(tensor[due].unsqueeze(0) for tensor in actual) |
| ref = tuple(tensor[due].unsqueeze(0) for tensor in ref) |
| exact = bench.exact_state(inputs)[due].unsqueeze(0) |
| passed, ratio = bench.judge(out, ref, exact) |
| assert passed, f"deployment accuracy failed: {ratio}" |
| for name in fields: |
| assert torch.equal(getattr(pool, name)[~due], before[name][~due]), f"non-due/slot-0 write: {name}" |
| for name in ("kbuf", "ubuf"): |
| assert torch.equal(getattr(pool, name), before[name]), f"update buffer changed: {name}" |
| |
| |
| assert torch.equal(pool.hcnt, before["hcnt"]), "flush changed fill counters" |
| assert bool((pool.w[due] == 0).all()) |
| if is_kda: |
| assert bool((pool.gbuf[due] == 1).all()) |
| assert bool((pool.pcum[due] == 0).all()) |
| else: |
| assert bool((pool.p[due] == 1).all()) |
| return float(ratio) |
|
|
| stream = torch.cuda.Stream() |
| inputs = bench.structured_inputs(ns, device, seed=91) |
| restore(inputs) |
| torch.cuda.synchronize() |
| before = snapshot() |
| torch.cuda.synchronize() |
| with torch.cuda.stream(stream): |
| op.flush_inplace(pool, idx, q0) |
| stream.synchronize() |
| ratios = [verify(inputs, before)] |
| before = snapshot() |
| op.flush_inplace(pool, torch.empty(0, dtype=torch.int32, device=device), q0) |
| torch.cuda.synchronize() |
| assert all(torch.equal(getattr(pool, name), before[name]) for name in fields), "empty idx is not a no-op" |
|
|
| restore(inputs) |
| torch.cuda.synchronize() |
| graph = torch.cuda.CUDAGraph() |
| with torch.cuda.graph(graph, stream=stream): |
| op.flush_inplace(pool, idx, q0) |
| for seed in (92, 93): |
| inputs = bench.structured_inputs(ns, device, seed=seed) |
| |
| generator = torch.Generator(device=device).manual_seed(seed) |
| inputs = (*inputs[:-1], torch.randn(128, 4, device=device, generator=generator)) |
| restore(inputs) |
| before = snapshot() |
| torch.cuda.synchronize() |
| graph.replay() |
| torch.cuda.synchronize() |
| ratios.append(verify(inputs, before)) |
| |
| pool.hcnt[due] = 0 |
| after = snapshot() |
| torch.cuda.synchronize() |
| graph.replay() |
| torch.cuda.synchronize() |
| assert all(torch.equal(getattr(pool, name), after[name]) for name in fields), "replay with no due heads is not a no-op" |
| print(json.dumps({ |
| "status": "pass", "gpu": torch.cuda.get_device_name(), |
| "compute_capability": list(torch.cuda.get_device_capability()), |
| "checks": ["runtime digests", "HF get_local_kernel", "mixed due heads", "unselected slots", "reserved slot 0", "window resets", "fill counters unchanged", "empty idx", "non-default stream", "graph capture", "graph replay with changed inputs and q0", "no-op replay with no due heads"], |
| "deployment_worst_head_error_ratios": ratios, |
| }, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
