"""Profile exact SageAttention2 preparation, mainloop, and tail scheduling on real H3 tensors.""" from __future__ import annotations import argparse import hashlib import json import math from pathlib import Path import torch from h3_blackwell_runtime.sage2_entry import prepare_v from profile_attention_path import prepare_qkv, representative_attention_inputs, summarize CTA_Q = 128 CTA_K = 64 WARP_Q = 32 WARP_K = 64 V_SCALE_MAX = 2.25 def tensor_sha256(value: torch.Tensor) -> str: host_bytes = value.detach().contiguous().view(torch.uint8).cpu().numpy() return hashlib.sha256(memoryview(host_bytes)).hexdigest() def event_measure(fn, *, warmup: int, iterations: int): for _ in range(warmup): fn() torch.cuda.synchronize() values = [] result = None for _ in range(iterations): started = torch.cuda.Event(enable_timing=True) finished = torch.cuda.Event(enable_timing=True) started.record() result = fn() finished.record() finished.synchronize() values.append(started.elapsed_time(finished) / 1000.0) return summarize(values), result def quantize_qk(q: torch.Tensor, k: torch.Tensor, km: torch.Tensor): import sageattention.core as sage_core return sage_core.per_warp_int8_cuda( q, k, km, BLKQ=CTA_Q, WARPQ=WARP_Q, BLKK=CTA_K, tensor_layout="NHD", ) def quantize_v(v: torch.Tensor): import sageattention.core as sage_core return sage_core.per_channel_fp8( v, tensor_layout="NHD", scale_max=V_SCALE_MAX, smooth_v=False, ) def run_mainloop( q_int8: torch.Tensor, k_int8: torch.Tensor, v_fp8: torch.Tensor, q_scale: torch.Tensor, k_scale: torch.Tensor, v_scale: torch.Tensor, output: torch.Tensor, ) -> torch.Tensor: import sageattention.core as sage_core sage_core.sm89_compile.qk_int8_sv_f8_accum_f16_fuse_v_scale_attn_inst_buf( q_int8, k_int8, v_fp8, output, q_scale, k_scale, v_scale, 0, 0, 2, output.shape[-1] ** -0.5, 0, ) return output def prepare_quantized(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor): km = k.mean(dim=1, keepdim=True) q_int8, q_scale, k_int8, k_scale = quantize_qk(q, k, km) v_fp8, v_scale, _ = quantize_v(v) output = torch.empty(q.shape, dtype=q.dtype, device=q.device) return km, q_int8, q_scale, k_int8, k_scale, v_fp8, v_scale, output def tail_row( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, *, warmup: int, iterations: int, ) -> dict: quantized = prepare_quantized(q, k, v) _, q_int8, q_scale, k_int8, k_scale, v_fp8, v_scale, output = quantized timing, result = event_measure( lambda: run_mainloop( q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale, output, ), warmup=warmup, iterations=iterations, ) q_len = q.shape[1] kv_len = k.shape[1] q_ctas = math.ceil(q_len / CTA_Q) k_iterations = math.ceil(kv_len / CTA_K) return { "q_len": q_len, "kv_len": kv_len, "q_ctas_per_head": q_ctas, "k_iterations_per_cta": k_iterations, "q_tail_rows": q_len % CTA_Q, "k_tail_rows": kv_len % CTA_K, "scheduled_q_rows": q_ctas * CTA_Q, "q_row_efficiency": q_len / (q_ctas * CTA_Q), "mainloop": timing, "checksum": result.float().sum().item(), } def difference(actual: torch.Tensor, expected: torch.Tensor) -> dict: delta = actual.float() - expected.float() return { "equal": torch.equal(actual, expected), "max_abs": delta.abs().max().item(), "mean_abs": delta.abs().mean().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("--attention", default="sage2", choices=("sage2",)) parser.add_argument("--warmup", type=int, default=3) parser.add_argument("--iterations", type=int, default=10) parser.add_argument("--tail-iterations", type=int, default=5) parser.add_argument("--skip-tail-study", action="store_true") parser.add_argument("--cuda-profiler-capture", action="store_true") parser.add_argument("--expected-output-sha256") 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) with torch.inference_mode(): q, k, v, _ = prepare_qkv(block, hidden, rotation, None) km_timing, km = event_measure( lambda: k.mean(dim=1, keepdim=True), warmup=args.warmup, iterations=args.iterations, ) qk_timing, qk = event_measure( lambda: quantize_qk(q, k, km), warmup=args.warmup, iterations=args.iterations, ) q_int8, q_scale, k_int8, k_scale = qk import sageattention.core as sage_core import sageattention.quant as sage_quant q_int8_probe = torch.empty(q.shape, dtype=torch.int8, device=q.device) q_scale_probe = torch.empty_like(q_scale) q_quant_timing, _ = event_measure( lambda: sage_quant._fused.quant_per_warp_int8_cuda( q, q_int8_probe, q_scale_probe, CTA_Q, WARP_Q, 0, ), warmup=args.warmup, iterations=args.iterations, ) k_int8_probe = torch.empty(k.shape, dtype=torch.int8, device=k.device) k_scale_probe = torch.empty_like(k_scale) k_quant_timing, _ = event_measure( lambda: sage_quant._fused.quant_per_block_int8_fuse_sub_mean_cuda( k, km.squeeze(1), k_int8_probe, k_scale_probe, CTA_K, 0, ), warmup=args.warmup, iterations=args.iterations, ) padded_k = math.ceil(v.shape[1] / CTA_K) * CTA_K v_transposed = torch.empty( (v.shape[0], v.shape[3], v.shape[2], padded_k), dtype=v.dtype, device=v.device, ) v_transpose_timing, _ = event_measure( lambda: sage_quant._fused.transpose_pad_permute_cuda(v, v_transposed, 0), warmup=args.warmup, iterations=args.iterations, ) v_fp8_probe = torch.empty_like(v_transposed, dtype=torch.float8_e4m3fn) v_scale_probe = torch.empty( (v.shape[0], v.shape[2], v.shape[3]), dtype=torch.float32, device=v.device, ) v_scale_quant_timing, _ = event_measure( lambda: sage_quant._fused.scale_fuse_quant_cuda( v_transposed, v_fp8_probe, v_scale_probe, v.shape[1], V_SCALE_MAX, 0, ), warmup=args.warmup, iterations=args.iterations, ) v_timing, vq = event_measure( lambda: quantize_v(v), warmup=args.warmup, iterations=args.iterations, ) v_fp8, v_scale, _ = vq candidate_v_timing, candidate_vq = event_measure( lambda: prepare_v(v, scale_max=V_SCALE_MAX), warmup=args.warmup, iterations=args.iterations, ) candidate_v_fp8, candidate_v_scale = candidate_vq output = torch.empty(q.shape, dtype=q.dtype, device=q.device) mainloop_fn = lambda: run_mainloop( q_int8, k_int8, v_fp8, q_scale, k_scale, v_scale, output, ) if args.cuda_profiler_capture: for _ in range(args.warmup): mainloop_fn() torch.cuda.synchronize() torch.cuda.cudart().cudaProfilerStart() captured = mainloop_fn() torch.cuda.synchronize() torch.cuda.cudart().cudaProfilerStop() report = { "metadata": metadata, "capture": "one unchanged prequantized Sage2 mainloop", "q_shape": list(q.shape), "k_shape": list(k.shape), "v_shape": list(v.shape), "checksum": captured.float().sum().item(), "scheduler": { "cta_q": CTA_Q, "cta_k": CTA_K, "warp_q": WARP_Q, "warp_k": WARP_K, "warps_per_cta": 4, "threads_per_cta": 128, "dynamic_shared_memory_bytes": 32768, "q_ctas_per_head": math.ceil(q.shape[1] / CTA_Q), "heads": q.shape[2], "grid_ctas": math.ceil(q.shape[1] / CTA_Q) * q.shape[2] * q.shape[0], "k_iterations_per_cta": math.ceil(k.shape[1] / CTA_K), "explicit_pipeline_stages": 2, }, } 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) return mainloop_timing, manual_output = event_measure( mainloop_fn, warmup=args.warmup, iterations=args.iterations, ) candidate_output = torch.empty_like(output) run_mainloop( q_int8, k_int8, candidate_v_fp8, q_scale, k_scale, candidate_v_scale, candidate_output, ) torch.cuda.synchronize() reference = __import__("sageattention").sageattn( q, k, v, tensor_layout="NHD", is_causal=False, smooth_k=False, ) torch.cuda.synchronize() output_sha256 = tensor_sha256(manual_output) expected_output_matches = ( args.expected_output_sha256 is None or output_sha256 == args.expected_output_sha256 ) tail_study = [] if not args.skip_tail_study: q_lengths = sorted({ (q.shape[1] // CTA_Q) * CTA_Q, (q.shape[1] // CTA_Q) * CTA_Q + 1, q.shape[1], }) kv_lengths = sorted({ (k.shape[1] // CTA_K) * CTA_K, (k.shape[1] // CTA_K) * CTA_K + 1, k.shape[1], }) for q_len in q_lengths: tail_study.append({ "sweep": "q_tail_fixed_kv", **tail_row( q[:, :q_len], k, v, warmup=args.warmup, iterations=args.tail_iterations, ), }) for kv_len in kv_lengths: tail_study.append({ "sweep": "kv_tail_fixed_q", **tail_row( q, k[:, :kv_len], v[:, :kv_len], warmup=args.warmup, iterations=args.tail_iterations, ), }) report = { "metadata": metadata, "q_shape": list(q.shape), "k_shape": list(k.shape), "v_shape": list(v.shape), "warmup": args.warmup, "iterations": args.iterations, "phase_timings": { "k_mean_and_smoothing_preparation": km_timing, "qk_int8_quantization": qk_timing, "q_int8_quantization": q_quant_timing, "k_int8_subtract_mean_quantization": k_quant_timing, "v_fp8_transpose_scale_quantization": v_timing, "v_transpose_pad_permute": v_transpose_timing, "v_scale_fp8_quantization": v_scale_quant_timing, "vortex_direct_v_fp8_preparation": candidate_v_timing, "fused_mainloop": mainloop_timing, }, "fused_mainloop_phases": { "int8_qk": "fused inside qk_int_sv_f8_attn_kernel", "scale_application": "fused inside qk_int_sv_f8_attn_kernel", "online_softmax": "fused inside qk_int_sv_f8_attn_kernel", "pv_accumulation": "fused inside qk_int_sv_f8_attn_kernel", "final_normalization_and_output": "fused inside qk_int_sv_f8_attn_kernel", "timing_policy": "Do not assign independent wall time without changing the exact kernel schedule; use source-correlated hardware counters.", }, "scheduler": { "cta_q": CTA_Q, "cta_k": CTA_K, "warp_q": WARP_Q, "warp_k": WARP_K, "warps_per_cta": 4, "threads_per_cta": 128, "dynamic_shared_memory_bytes": 32768, "q_ctas_per_head": math.ceil(q.shape[1] / CTA_Q), "heads": q.shape[2], "grid_ctas": math.ceil(q.shape[1] / CTA_Q) * q.shape[2] * q.shape[0], "k_iterations_per_cta": math.ceil(k.shape[1] / CTA_K), "q_tail_rows": q.shape[1] % CTA_Q, "k_tail_rows": k.shape[1] % CTA_K, "explicit_pipeline_stages": 2, }, "manual_decomposition_vs_public_sage2": difference(manual_output, reference), "vortex_v_fp8_vs_sage2": difference(candidate_v_fp8, v_fp8), "vortex_v_scale_vs_sage2": difference(candidate_v_scale, v_scale), "vortex_v_mainloop_vs_sage2": difference(candidate_output, manual_output), "manual_checksum": manual_output.float().sum().item(), "reference_checksum": reference.float().sum().item(), "output_sha256": output_sha256, "expected_output_sha256": args.expected_output_sha256, "expected_output_matches": expected_output_matches, "tail_study": tail_study, } 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 not expected_output_matches: raise RuntimeError( f"output SHA256 mismatch: expected {args.expected_output_sha256}, got {output_sha256}" ) if __name__ == "__main__": main()