133 lines
4.6 KiB
Python
133 lines
4.6 KiB
Python
"""Capture fixed aligned prequantized Sage2 fixtures for VEA-B Phase 2B."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
import platform
|
|
from pathlib import Path
|
|
|
|
import torch
|
|
|
|
|
|
SHAPE = {"batch": 1, "heads": 1, "q": 128, "kv": 192, "d": 128}
|
|
SEED = 73021
|
|
|
|
|
|
def tensor_sha256(value: torch.Tensor) -> str:
|
|
data = value.detach().contiguous().view(torch.uint8).cpu().numpy()
|
|
return hashlib.sha256(memoryview(data)).hexdigest()
|
|
|
|
|
|
def file_sha256(path: Path) -> str:
|
|
digest = hashlib.sha256()
|
|
with path.open("rb") as handle:
|
|
for chunk in iter(lambda: handle.read(16 * 1024 * 1024), b""):
|
|
digest.update(chunk)
|
|
return digest.hexdigest()
|
|
|
|
|
|
def save(output_dir: Path, name: str, value: torch.Tensor) -> dict:
|
|
host = value.detach().contiguous().cpu()
|
|
path = output_dir / f"{name}.pt"
|
|
torch.save({"tensor": host}, path)
|
|
loaded = torch.load(path, map_location="cpu", weights_only=True)["tensor"]
|
|
if not torch.equal(host, loaded):
|
|
raise RuntimeError(f"reload mismatch for {name}")
|
|
return {
|
|
"path": str(path),
|
|
"shape": list(host.shape),
|
|
"dtype": str(host.dtype),
|
|
"tensor_sha256": tensor_sha256(host),
|
|
"file_sha256": file_sha256(path),
|
|
"size_bytes": path.stat().st_size,
|
|
"reload_verified": True,
|
|
}
|
|
|
|
|
|
def public_mainloop(q, k, v, q_scale, k_scale, v_scale):
|
|
import sageattention.core as sage_core
|
|
|
|
output = torch.empty(q.shape, dtype=torch.bfloat16, device=q.device)
|
|
sage_core.sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf(
|
|
q, k, v, output, q_scale, k_scale, v_scale,
|
|
0, 0, 2, SHAPE["d"] ** -0.5, 0,
|
|
)
|
|
return output
|
|
|
|
|
|
def main() -> None:
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
|
parser.add_argument("--image", required=True)
|
|
args = parser.parse_args()
|
|
args.output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
generator = torch.Generator(device="cuda").manual_seed(SEED)
|
|
q_bf16 = torch.randn((1, 128, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator)
|
|
k_bf16 = torch.randn((1, 192, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator)
|
|
v_bf16 = torch.randn((1, 192, 1, 128), dtype=torch.bfloat16, device="cuda", generator=generator)
|
|
|
|
import sageattention.core as sage_core
|
|
|
|
with torch.inference_mode():
|
|
k_mean = k_bf16.mean(dim=1, keepdim=True)
|
|
q_int8, q_scale, k_int8, k_scale = sage_core.per_warp_int8_cuda(
|
|
q_bf16, k_bf16, k_mean, BLKQ=128, WARPQ=32, BLKK=64,
|
|
tensor_layout="NHD",
|
|
)
|
|
v_fp8, v_scale, _ = sage_core.per_channel_fp8(
|
|
v_bf16, tensor_layout="NHD", scale_max=2.25, smooth_v=False,
|
|
)
|
|
public_output = public_mainloop(q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale)
|
|
torch.cuda.synchronize()
|
|
|
|
tensors = {
|
|
"q_bf16": q_bf16,
|
|
"k_bf16": k_bf16,
|
|
"v_bf16": v_bf16,
|
|
"k_mean_bf16": k_mean,
|
|
"q_int8": q_int8,
|
|
"q_scale_fp32": q_scale,
|
|
"k_int8": k_int8,
|
|
"k_scale_fp32": k_scale,
|
|
"v_fp8_e4m3": v_fp8,
|
|
"v_scale_fp32": v_scale,
|
|
"public_output_bf16": public_output,
|
|
}
|
|
records = {name: save(args.output_dir, name, value) for name, value in tensors.items()}
|
|
manifest = {
|
|
"schema": "vortex-exact-phase2b-aligned-fixtures",
|
|
"version": 1,
|
|
"status": "captured_and_reload_verified",
|
|
"shape": SHAPE,
|
|
"epochs": 3,
|
|
"seed": SEED,
|
|
"scope": "prepared BF16 retained only for provenance; prototype inputs are prequantized tensors",
|
|
"reference": {
|
|
"sageattention_version": "2.2.0",
|
|
"sageattention_commit": "d1a57a546c3d395b1ffcbeecc66d81db76f3b4b5",
|
|
"mainloop": "qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf",
|
|
"tensor_layout": "NHD",
|
|
"causal": False,
|
|
"qk_quant_granularity": "per_warp",
|
|
"v_scale_max": 2.25,
|
|
},
|
|
"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(),
|
|
},
|
|
"tensors": records,
|
|
}
|
|
manifest_path = args.output_dir / "manifest.json"
|
|
manifest_path.write_text(json.dumps(manifest, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps(manifest, indent=2), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|