268 lines
12 KiB
Python
268 lines
12 KiB
Python
|
"""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()
|