217 lines
8 KiB
Python
217 lines
8 KiB
Python
|
"""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()
|