154 lines
4.6 KiB
Python
154 lines
4.6 KiB
Python
"""Byte-compare isolated SM89 extensions on adversarial short NHD shapes."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import importlib.util
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
CASES = (
|
|
(1, 1),
|
|
(31, 63),
|
|
(32, 64),
|
|
(33, 65),
|
|
(63, 127),
|
|
(64, 128),
|
|
(65, 129),
|
|
(127, 63),
|
|
(128, 64),
|
|
(129, 65),
|
|
(191, 191),
|
|
(192, 192),
|
|
(193, 193),
|
|
)
|
|
|
|
|
|
def load_extension(path: Path):
|
|
spec = importlib.util.spec_from_file_location("sageattention._qattn_sm89", path)
|
|
if spec is None or spec.loader is None:
|
|
raise RuntimeError(f"cannot load extension from {path}")
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
return module
|
|
|
|
|
|
def digest(value: torch.Tensor) -> str:
|
|
data = value.detach().contiguous().view(torch.uint8).cpu().numpy()
|
|
return hashlib.sha256(memoryview(data)).hexdigest()
|
|
|
|
|
|
def launch(extension, q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale):
|
|
output = torch.empty(q_int8.shape, dtype=torch.bfloat16, device=q_int8.device)
|
|
extension.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
|
|
q_int8,
|
|
k_int8,
|
|
v_fp8,
|
|
output,
|
|
q_scale,
|
|
k_scale,
|
|
v_scale,
|
|
0,
|
|
0,
|
|
2,
|
|
128**-0.5,
|
|
0,
|
|
)
|
|
return output
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--baseline", type=Path, required=True)
|
|
parser.add_argument("--candidate", type=Path, required=True)
|
|
parser.add_argument("--output", type=Path)
|
|
parser.add_argument("--heads", type=int, default=2)
|
|
parser.add_argument("--seed", type=int, default=73021)
|
|
parser.add_argument("--quick", action="store_true", help="Run only one one-tile and one multi-tile case.")
|
|
args = parser.parse_args()
|
|
|
|
import sageattention.core as sage_core
|
|
|
|
baseline = load_extension(args.baseline)
|
|
candidate = load_extension(args.candidate)
|
|
cases = ((1, 1), (129, 193)) if args.quick else CASES
|
|
results = []
|
|
|
|
with torch.inference_mode():
|
|
for index, (q_len, kv_len) in enumerate(cases):
|
|
generator = torch.Generator(device="cuda").manual_seed(args.seed + index)
|
|
q = torch.randn(
|
|
(1, q_len, args.heads, 128),
|
|
dtype=torch.bfloat16,
|
|
device="cuda",
|
|
generator=generator,
|
|
)
|
|
k = torch.randn(
|
|
(1, kv_len, args.heads, 128),
|
|
dtype=torch.bfloat16,
|
|
device="cuda",
|
|
generator=generator,
|
|
)
|
|
v = torch.randn(
|
|
(1, kv_len, args.heads, 128),
|
|
dtype=torch.bfloat16,
|
|
device="cuda",
|
|
generator=generator,
|
|
)
|
|
k_mean = k.mean(dim=1, keepdim=True)
|
|
q_int8, q_scale, k_int8, k_scale = sage_core.per_warp_int8_cuda(
|
|
q,
|
|
k,
|
|
k_mean,
|
|
BLKQ=128,
|
|
WARPQ=32,
|
|
BLKK=64,
|
|
tensor_layout="NHD",
|
|
)
|
|
v_fp8, v_scale, _ = sage_core.per_channel_fp8(
|
|
v,
|
|
tensor_layout="NHD",
|
|
scale_max=2.25,
|
|
smooth_v=False,
|
|
)
|
|
expected = launch(
|
|
baseline, q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale,
|
|
)
|
|
actual = launch(
|
|
candidate, q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale,
|
|
)
|
|
torch.cuda.synchronize()
|
|
delta = actual.float() - expected.float()
|
|
result = {
|
|
"q_len": q_len,
|
|
"kv_len": kv_len,
|
|
"equal": torch.equal(actual, expected),
|
|
"different_elements": int(torch.count_nonzero(actual != expected).item()),
|
|
"max_abs": delta.abs().max().item(),
|
|
"baseline_sha256": digest(expected),
|
|
"candidate_sha256": digest(actual),
|
|
}
|
|
results.append(result)
|
|
print(json.dumps(result), flush=True)
|
|
if not result["equal"]:
|
|
raise RuntimeError(f"short-shape parity failed: {result}")
|
|
|
|
report = {
|
|
"device": torch.cuda.get_device_name(),
|
|
"baseline": str(args.baseline),
|
|
"candidate": str(args.candidate),
|
|
"heads": args.heads,
|
|
"seed": args.seed,
|
|
"cases": results,
|
|
}
|
|
if args.output is not None:
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|