File size: 2,486 Bytes
22a73d3 | 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 | #!/usr/bin/env python3
"""Measure the zero-remainder, unweighted shared-layout reblock path."""
from __future__ import annotations
import argparse
import json
import statistics
import torch
from h3_sparse_attention.processor import H3SparseAttentionConfig, PackedLayout, _Controller
from h3_sparse_attention.spark_integration import _landmark_tree_v2_combined_permutations
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--device", type=int, default=0)
parser.add_argument("--heads", type=int, default=24)
parser.add_argument("--repeats", type=int, default=20)
args = parser.parse_args()
torch.cuda.set_device(args.device)
device = torch.device("cuda", args.device)
grid = (21, 54, 64)
tokens = 21 * 54 * 64
ids = torch.arange(tokens, device=device)
positions = torch.stack(
torch.meshgrid(
*(torch.arange(size, device=device) for size in grid), indexing="ij"
), dim=-1,
).reshape(tokens, 3)
layout = PackedLayout(ids, ids.clone(), grid, tokens, tokens, positions)
torch.manual_seed(20261005)
query = torch.randn(
1, tokens, args.heads, 128, device=device, dtype=torch.bfloat16
)
key = torch.randn_like(query)
controller = _Controller(H3SparseAttentionConfig.spark(
20,
landmark_tree_v2_layout_reuse="q_from_k",
landmark_tree_v2_m2_side="none",
))
def call():
return _landmark_tree_v2_combined_permutations(
controller, query, key, layout, expand_shared=False
)
for _ in range(3):
call()
torch.cuda.synchronize()
values = []
for _ in range(args.repeats):
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)
start.record()
call()
end.record()
end.synchronize()
values.append(start.elapsed_time(end))
plans = [
value for value in controller.rope_sol_key_clustering_static.values()
if hasattr(value, "graph_active")
]
print(json.dumps({
"device": args.device,
"gpu": torch.cuda.get_device_name(device),
"shape": [1, tokens, args.heads, 128],
"mean_ms": statistics.mean(values),
"median_ms": statistics.median(values),
"min_ms": min(values),
"max_ms": max(values),
"graph_active": bool(plans and plans[0].graph_active),
}, indent=2))
if __name__ == "__main__":
main()
|