| |
| """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() |
|
|
