"""Attribute Sage2 and Sol-Attn latency on representative real H3 tensors.""" from __future__ import annotations import argparse import json import time from contextlib import contextmanager from pathlib import Path import torch from profile_attention_path import prepare_qkv, representative_attention_inputs, summarize, sync class StageHooks: def __init__(self): self.current: dict[str, float] = {} self.originals: list[tuple[object, str, object]] = [] def patch(self, owner: object, attribute: str, stage: str) -> None: original = getattr(owner, attribute) def wrapped(*args, **kwargs): sync() started = time.perf_counter() result = original(*args, **kwargs) sync() self.current[stage] = self.current.get(stage, 0.0) + time.perf_counter() - started return result self.originals.append((owner, attribute, original)) setattr(owner, attribute, wrapped) def restore(self) -> None: for owner, attribute, original in reversed(self.originals): setattr(owner, attribute, original) self.originals.clear() @contextmanager def installed_hooks(configure): hooks = StageHooks() try: configure(hooks) yield hooks finally: hooks.restore() def measure(fn, configure_hooks, *, warmup: int, iterations: int) -> tuple[dict, torch.Tensor]: with installed_hooks(configure_hooks) as hooks, torch.inference_mode(): for _ in range(warmup): hooks.current = {} fn() rows = [] output = None for _ in range(iterations): hooks.current = {} sync() started = time.perf_counter() output = fn() sync() row = dict(hooks.current) row["total"] = time.perf_counter() - started rows.append(row) names = sorted({name for row in rows for name in row}) return {name: summarize([row.get(name, 0.0) for row in rows]) for name in names}, output def add_residual(timings: dict, total: str, children: tuple[str, ...], name: str) -> None: values = [] count = timings[total]["count"] # The summary alone cannot reconstruct paired iterations. This residual is # therefore an explicitly labeled median estimate, not a distribution. estimate = timings[total]["p50_s"] - sum(timings[child]["p50_s"] for child in children if child in timings) values.extend([max(estimate, 0.0)] * count) timings[name] = summarize(values) def difference(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float]: delta = actual.float() - expected.float() return { "max_abs": delta.abs().max().item(), "mean_abs": delta.abs().mean().item(), "relative_l2": (delta.norm() / expected.float().norm().clamp_min(1e-12)).item(), } def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") parser.add_argument("--output", type=Path, required=True) parser.add_argument("--width", type=int, default=1344) parser.add_argument("--height", type=int, default=768) parser.add_argument("--frames", type=int, default=124) parser.add_argument("--steps", type=int, default=12) parser.add_argument("--sampler-step", type=int, default=1) parser.add_argument("--seed", type=int, default=440420) parser.add_argument("--text-tokens", type=int, default=100) parser.add_argument("--block-index", type=int, default=24) parser.add_argument("--sol-tau", type=float, default=1.3) parser.add_argument("--warmup", type=int, default=3) parser.add_argument("--iterations", type=int, default=10) parser.add_argument("--attention", default="sage2", choices=("sage2",)) parser.add_argument("--sage-only", action="store_true") parser.add_argument("--device", default="cuda") return parser.parse_args() def main() -> None: args = parse_args() block, hidden, rotation, segments, metadata = representative_attention_inputs(args) q, k, v, _ = prepare_qkv(block, hidden, rotation, None) sequence = q.shape[1] conditioning_stop = segments[-1][0] sink_blocks = (0, (conditioning_stop + 63) // 64) import sageattention.core as sage_core def configure_sage(hooks: StageHooks) -> None: hooks.patch(sage_core, "per_warp_int8_cuda", "qk_quantize") hooks.patch(sage_core, "per_channel_fp8", "v_quantize") hooks.patch( sage_core.sm89_compile, "qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf", "attention_forward", ) sage_timings, sage_output = measure( lambda: sage_core.sageattn( q, k, v, tensor_layout="NHD", is_causal=False, smooth_k=False, ), configure_sage, warmup=args.warmup, iterations=args.iterations, ) add_residual( sage_timings, "total", ("qk_quantize", "v_quantize", "attention_forward"), "framework_and_k_smoothing_estimate", ) sol_results = [] if not args.sage_only: import sol_kernel.fwd as sol_fwd import sol_kernel.preprocess as sol_preprocess import sol_kernel.quant as sol_quant def configure_sol(hooks: StageHooks) -> None: hooks.patch(sol_preprocess, "reduce_quantize_k", "k_summary_and_optional_quantize") hooks.patch(sol_preprocess, "_reduce_v", "v_summary") hooks.patch(sol_preprocess, "_compute_diag_threshold", "routing_threshold") hooks.patch(sol_preprocess, "quantize_q_with_threshold", "q_quantize_and_threshold") hooks.patch(sol_quant, "quantize_v_per_channel", "v_int8_quantize") hooks.patch(sol_fwd, "prepare", "prepare_total") for int8_qk, int8_pv in ((False, False), (True, False), (True, True)): sol_timings, sol_output = measure( lambda int8_qk=int8_qk, int8_pv=int8_pv: sol_fwd.sol_attn( q, k, v, tau=args.sol_tau, thresh_type="diag", int8_qk=int8_qk, int8_pv=int8_pv, sink_blocks=sink_blocks, ), configure_sol, warmup=args.warmup, iterations=args.iterations, ) add_residual(sol_timings, "total", ("prepare_total",), "routed_forward_estimate") prepare_children = ( "k_summary_and_optional_quantize", "v_summary", "routing_threshold", "q_quantize_and_threshold", "v_int8_quantize", ) add_residual(sol_timings, "prepare_total", prepare_children, "prepare_other_estimate") sol_results.append({ "int8_qk": int8_qk, "int8_pv": int8_pv, "tau": args.sol_tau, "sink_blocks": list(sink_blocks), "timings": sol_timings, "difference_vs_sage2": difference(sol_output, sage_output), }) report = { "metadata": metadata, "sequence": sequence, "conditioning_stop": conditioning_stop, "warmup": args.warmup, "iterations": args.iterations, "measurement_policy": { "component_timings": "CUDA-synchronized nested probes; attribution only", "residual_timings": "difference of component p50 values; labeled estimates", "end_to_end": "must be measured separately without component synchronization", }, "sage2": { "version": "2.2.0", "timings": sage_timings, "checksum": sage_output.float().sum().item(), }, "sol": sol_results, } 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 __name__ == "__main__": main()