"""Build and validate the fixed aligned VEA-B Phase 2B numerical prototype.""" from __future__ import annotations import argparse import contextlib import hashlib import io import json import math import os import platform import statistics import sys import time from pathlib import Path import torch from torch.utils.cpp_extension import load PROJECT = Path(__file__).resolve().parents[1] KERNELS = PROJECT / "kernels" NAMES = ("output", "qk_scores", "m", "d", "probability_fp8", "rescale", "pv_fp16", "ro_fp32", "reciprocal", "clocks") ATTRIBUTE_NAMES = ("registers_per_thread", "static_shared_bytes", "local_bytes_per_thread", "max_threads_per_block", "binary_version", "ptx_version", "max_dynamic_shared_bytes", "resident_ctas_per_sm") KERNEL_NAMES = ("vea_b_fast", "vea_b_capture", "sage_checkpoint") def tensor_sha256(value: torch.Tensor) -> str: data = value.detach().contiguous().view(torch.uint8).cpu().numpy() return hashlib.sha256(memoryview(data)).hexdigest() def percentile(values: list[float], fraction: float) -> float: ordered = sorted(values) rank = (len(ordered) - 1) * fraction low = int(rank) high = min(low + 1, len(ordered) - 1) weight = rank - low return ordered[low] * (1 - weight) + ordered[high] * weight def summarize(values: list[float]) -> dict: mean = statistics.fmean(values) return { "count": len(values), "p10_ms": percentile(values, 0.10), "p50_ms": percentile(values, 0.50), "p90_ms": percentile(values, 0.90), "p95_ms": percentile(values, 0.95), "mean_ms": mean, "stdev_ms": statistics.pstdev(values), "coefficient_of_variation": statistics.pstdev(values) / mean if mean else 0.0, "min_ms": min(values), "max_ms": max(values), } def benchmark(fn, warmup: int, iterations: int, batch: int) -> dict: for _ in range(warmup): fn() torch.cuda.synchronize() values = [] for _ in range(iterations): start = torch.cuda.Event(enable_timing=True) end = torch.cuda.Event(enable_timing=True) start.record() for _ in range(batch): fn() end.record() end.synchronize() values.append(start.elapsed_time(end) / batch) result = summarize(values) result["launches_per_sample"] = batch return result def build_extension(build_dir: Path, build_log: Path): build_dir.mkdir(parents=True, exist_ok=True) stream = io.StringIO() started = time.perf_counter() with contextlib.redirect_stdout(stream), contextlib.redirect_stderr(stream): extension = load( name="vortex_exact_phase2b_numeric", sources=[str(KERNELS / "vea_b_numeric.cpp"), str(KERNELS / "vea_b_numeric.cu")], build_directory=str(build_dir), extra_cflags=["-O3"], extra_cuda_cflags=["-O3", "-lineinfo", "--ptxas-options=-v", "--maxrregcount=200", "-std=c++17"], verbose=True, ) text = stream.getvalue() sys.stderr.write(text) build_log.parent.mkdir(parents=True, exist_ok=True) build_log.write_text(text, encoding="utf-8") return extension, time.perf_counter() - started def load_fixture(directory: Path) -> tuple[dict, dict[str, torch.Tensor]]: manifest = json.loads((directory / "manifest.json").read_text(encoding="utf-8")) tensors = {} for name, record in manifest["tensors"].items(): value = torch.load(directory / Path(record["path"]).name, map_location="cpu", weights_only=True)["tensor"] if tensor_sha256(value) != record["tensor_sha256"]: raise RuntimeError(f"fixture hash mismatch: {name}") tensors[name] = value.cuda() return manifest, tensors def public_launch(tensors: dict[str, torch.Tensor], output: torch.Tensor) -> torch.Tensor: import sageattention.core as sage_core sage_core.sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf( tensors["q_int8"], tensors["k_int8"], tensors["v_fp8_e4m3"], output, tensors["q_scale_fp32"], tensors["k_scale_fp32"], tensors["v_scale_fp32"], 0, 0, 2, 128 ** -0.5, 0, ) return output def comparison(actual: torch.Tensor, expected: torch.Tensor, name: str) -> dict: equal = torch.equal(actual, expected) result = { "name": name, "equal": equal, "actual_sha256": tensor_sha256(actual), "expected_sha256": tensor_sha256(expected), "different_elements": int(torch.count_nonzero(actual != expected).item()), } if not equal: mismatch = torch.nonzero(actual != expected, as_tuple=False)[0].cpu().tolist() index = tuple(mismatch) result["first_mismatch_index"] = mismatch result["first_divergent_epoch"] = mismatch[0] if actual.ndim and actual.shape[0] == 3 else None result["actual_first"] = actual[index].item() result["expected_first"] = expected[index].item() if actual.is_floating_point(): delta = actual.float() - expected.float() result["max_abs"] = float(delta.abs().max().item()) return result def analyze_overlap(clocks: torch.Tensor) -> dict: host = clocks.cpu() transitions = [] for epoch in range(2): pv_start = int(host[epoch, 4:8, 0].min().item()) pv_end = int(host[epoch, 4:8, 1].max().item()) qk_start = int(host[epoch + 1, 0:4, 0].min().item()) qk_end = int(host[epoch + 1, 0:4, 1].max().item()) overlap = max(0, min(pv_end, qk_end) - max(pv_start, qk_start)) transitions.append({"pv_epoch": epoch, "qk_epoch": epoch + 1, "overlap_clocks": overlap, "positive": overlap > 0}) return {"transitions": transitions, "all_positive": all(item["positive"] for item in transitions)} def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--fixture-dir", type=Path, required=True) parser.add_argument("--output", type=Path, required=True) parser.add_argument("--build-dir", type=Path, required=True) parser.add_argument("--build-log", type=Path, required=True) parser.add_argument("--image", required=True) parser.add_argument("--warmup", type=int, default=100) parser.add_argument("--iterations", type=int, default=500) parser.add_argument("--repeat", type=int, default=1000) parser.add_argument("--timing-batch", type=int, default=1) parser.add_argument("--capture", choices=("none", "vea", "sage_checkpoint"), default="none") parser.add_argument("--cuda-profiler-capture", action="store_true") return parser.parse_args() def main() -> None: args = parse_args() extension, build_seconds = build_extension(args.build_dir, args.build_log) manifest, tensors = load_fixture(args.fixture_dir) inputs = ( tensors["q_int8"], tensors["k_int8"], tensors["v_fp8_e4m3"], tensors["q_scale_fp32"], tensors["k_scale_fp32"], tensors["v_scale_fp32"], ) attributes_tensor = extension.attributes() attributes = { kernel: {field: int(value) for field, value in zip(ATTRIBUTE_NAMES, row.tolist(), strict=True)} for kernel, row in zip(KERNEL_NAMES, attributes_tensor, strict=True) } if args.capture != "none": if args.cuda_profiler_capture: torch.cuda.cudart().cudaProfilerStart() if args.capture == "vea": output = torch.empty_like(tensors["public_output_bf16"]) captured = extension.vea_b_into(*inputs, output) else: captured = extension.sage_checkpoint(*inputs)[0] torch.cuda.synchronize() if args.cuda_profiler_capture: torch.cuda.cudart().cudaProfilerStop() report = { "status": "capture_only", "capture": args.capture, "output_sha256": tensor_sha256(captured), "attributes": attributes, } else: reference = extension.sage_checkpoint(*inputs) actual = extension.vea_b(*inputs, True) torch.cuda.synchronize() public_expected = tensors["public_output_bf16"] parity = [comparison(reference[0], public_expected, "sage_checkpoint_vs_public_output")] parity.extend(comparison(actual[index], reference[index], name) for index, name in enumerate(NAMES[:-1])) first_failure = next((item for item in parity if not item["equal"]), None) repeat_output = torch.empty_like(public_expected) for _ in range(args.repeat): extension.vea_b_into(*inputs, repeat_output) torch.cuda.synchronize() repeat_parity = comparison(repeat_output, public_expected, "repeated_vea_output") vea_output = torch.empty_like(public_expected) public_output = torch.empty_like(public_expected) vea_timing = benchmark( lambda: extension.vea_b_into(*inputs, vea_output), args.warmup, args.iterations, args.timing_batch) public_timing = benchmark( lambda: public_launch(tensors, public_output), args.warmup, args.iterations, args.timing_batch) torch.cuda.synchronize() report = { "status": "passed" if first_failure is None and repeat_parity["equal"] else "failed", "parity": parity, "first_failure": first_failure, "repeat": {"iterations": args.repeat, "parity": repeat_parity}, "overlap": analyze_overlap(actual[9]), "timing": { "vea_b": vea_timing, "public_sage2": public_timing, "prototype_speedup": public_timing["p50_ms"] / vea_timing["p50_ms"], }, "attributes": attributes, "advancement": { "byte_exact_final": all(item["equal"] for item in parity if item["name"] in ("output", "sage_checkpoint_vs_public_output")), "all_intermediates_exact": all(item["equal"] for item in parity), "no_deadlock_repeat": repeat_parity["equal"], "registers_at_most_200": attributes["vea_b_fast"]["registers_per_thread"] <= 200, "local_bytes_zero": attributes["vea_b_fast"]["local_bytes_per_thread"] == 0, "one_ten_warp_cta_resident": attributes["vea_b_fast"]["resident_ctas_per_sm"] >= 1, "qk_pv_overlap": analyze_overlap(actual[9])["all_positive"], "stable_timing": vea_timing["coefficient_of_variation"] <= 0.05, }, } report.update({ "schema": "vortex-exact-phase2b-aligned-prototype", "version": 1, "build_seconds": build_seconds, "fixture_manifest": str(args.fixture_dir / "manifest.json"), "environment": { "image": args.image, "gpu": torch.cuda.get_device_name(), "compute_capability": list(torch.cuda.get_device_capability()), "torch": torch.__version__, "cuda": torch.version.cuda, "python": platform.python_version(), "argv": sys.argv, "environment_switches": {key: value for key, value in os.environ.items() if key.startswith(("CUDA_", "TORCH_", "MAX_JOBS"))}, }, "sources": { "cpp_sha256": hashlib.sha256((KERNELS / "vea_b_numeric.cpp").read_bytes()).hexdigest(), "cuda_sha256": hashlib.sha256((KERNELS / "vea_b_numeric.cu").read_bytes()).hexdigest(), }, "claims": {"fixed_aligned_shape_only": True, "production_changed": False, "h3_integration": False}, }) args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8") print(json.dumps(report, indent=2), flush=True) if report.get("status") == "failed": raise SystemExit(2) if __name__ == "__main__": main()