File size: 5,753 Bytes
f700360 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 | # SPDX-License-Identifier: Apache-2.0
"""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:])
# Each selected (slot, head) is a separate checked head, never averaged across slots.
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}"
# The serving contract leaves hcnt at L ("flushed, empty"); the decode
# step consumes that sentinel. Clearing it in flush would race scanners.
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)
# structured_inputs always uses the same q0; vary it explicitly here.
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))
# Simulate the caller advancing the counters so no head is due.
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()
|