"""Gate the fused Sage2 entry path on captured complete H3 blocks.""" from __future__ import annotations import argparse import json import os import time import types from pathlib import Path import torch from h3_blackwell_runtime.checkpoint import H3Checkpoint from h3_blackwell_runtime.denoiser import H3PackedDenoiser from h3_blackwell_runtime.packing import H3PromptPacker from h3_blackwell_runtime.rope import h3_rope_rotation from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, beta_sigmas from h3_blackwell_runtime.t2v import random_av_latents from profile_h3_block import profile_module_kernels, summarize def compare(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float | int | bool]: different = torch.count_nonzero(actual != expected).item() delta = (actual.float() - expected.float()).abs() return { "equal": different == 0, "different_elements": different, "max_abs": float(delta.max()), "mean_abs": float(delta.mean()), } def timed(fn) -> tuple[float, torch.Tensor]: torch.cuda.synchronize() started = time.perf_counter() output = fn() torch.cuda.synchronize() return time.perf_counter() - started, output def memory(fn) -> dict[str, int]: torch.cuda.synchronize() torch.cuda.reset_peak_memory_stats() allocated = torch.cuda.memory_allocated() reserved = torch.cuda.memory_reserved() output = fn() torch.cuda.synchronize() result = { "allocated_bytes_before": allocated, "reserved_bytes_before": reserved, "peak_allocated_bytes": torch.cuda.max_memory_allocated(), "peak_reserved_bytes": torch.cuda.max_memory_reserved(), "peak_allocated_delta_bytes": torch.cuda.max_memory_allocated() - allocated, "peak_reserved_delta_bytes": torch.cuda.max_memory_reserved() - reserved, } del output return result def main() -> None: parser = argparse.ArgumentParser() parser.add_argument("--capture", type=Path) parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") parser.add_argument("--blocks", type=int, nargs="+", default=[0, 24, 49]) 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("--warmup", type=int, default=2) parser.add_argument("--rounds", type=int, default=8) parser.add_argument("--cuda-profiler-capture", action="store_true") parser.add_argument("--profile-mode", choices=("baseline", "candidate"), default="baseline") parser.add_argument( "--output", type=Path, default=Path("/output/h3-blackwell-runtime/benchmarks/gb10-sage2-p1-block-gate.json"), ) args = parser.parse_args() os.environ["H3_SAGE_QKV_LAYOUT"] = "strided_nhd" os.environ["H3_SAGE_ENTRY_FUSION"] = "0" torch.manual_seed(args.seed) checkpoint = H3Checkpoint(args.model) model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval() if args.capture is not None: captured = torch.load(args.capture / "input.pt", map_location="cuda", weights_only=False) hidden = captured["hidden"].to("cuda") timesteps = captured["timesteps"].to("cuda") positions = captured["position_ids"].to("cuda") segments = captured["segments"] workload = {"capture": str(args.capture)} else: packer = H3PromptPacker(checkpoint) video, audio, aligned_frames = random_av_latents( args.width, args.height, args.frames, args.seed, device="cuda", ) sigma = beta_sigmas(args.steps, device="cuda")[args.sampler_step - 1] native_audio = audio.to(torch.bfloat16) * (_audio_sigma(sigma) / sigma) text = torch.randn(1, args.text_tokens, 5376, device="cuda", dtype=torch.bfloat16) hidden, timesteps, segments, positions, _, _ = packer( text, video, native_audio, _model_sigma(sigma), ) workload = { "resolution": [args.width, args.height], "frames": aligned_frames, "steps": args.steps, "sampler_step": args.sampler_step, "seed": args.seed, "text_tokens": args.text_tokens, } rotation = h3_rope_rotation(positions.to("cuda"), model.backbone.inv_freq, torch.bfloat16) workload["tokens"] = hidden.shape[0] workload["hidden_shape"] = list(hidden.shape) workload["rotation_shape"] = list(rotation.shape) print(workload, flush=True) block_inputs = {} block_outputs = {} rows = [] with torch.inference_mode(): for index, (block, adaln) in enumerate(zip(model.backbone.blocks, model.backbone.adaln, strict=True)): if index in args.blocks: block_inputs[index] = hidden hidden = block(hidden, rotation, *adaln(timesteps), segments) if index in args.blocks: block_outputs[index] = hidden if index >= max(args.blocks): break baseline_forwards = {index: model.backbone.blocks[index].attention.forward for index in args.blocks} def fused_attention_forward( attention, x, rope_rotation, sequence_parallel=None, tensor_parallel=None, modulation=None, ): from h3_blackwell_runtime.sage2_entry import attention_nhd if sequence_parallel is not None or tensor_parallel is not None: raise ValueError("the Sage2 entry-fusion benchmark is single-GPU only") sequence = x.shape[0] inner = attention.heads * attention.head_dim qkv = attention.qkv_proj.forward_modulated(x, *modulation) if modulation is not None else attention.qkv_proj(x) q, k, v = qkv.split(inner, dim=-1) q = q.view(1, sequence, attention.heads, attention.head_dim) k = k.view(1, sequence, attention.heads, attention.head_dim) v = v.view(1, sequence, attention.heads, attention.head_dim) output = attention_nhd( q, k, v, rope_rotation, attention.q_norm_weight, attention.k_norm_weight, attention.eps, ) return attention.out_proj(output.reshape(sequence, inner)) candidate_forwards = { index: types.MethodType(fused_attention_forward, model.backbone.blocks[index].attention) for index in args.blocks } def set_fused(index: int, fused: bool) -> None: model.backbone.blocks[index].attention.forward = ( candidate_forwards[index] if fused else baseline_forwards[index] ) if args.cuda_profiler_capture: checksums = {} fused = args.profile_mode == "candidate" os.environ["H3_SAGE_ENTRY_FUSION"] = "1" if fused else "0" for index in args.blocks: block = model.backbone.blocks[index] set_fused(index, fused) adaln_values = model.backbone.adaln[index](timesteps) for _ in range(args.warmup): block(block_inputs[index], rotation, *adaln_values, segments) torch.cuda.synchronize() torch.cuda.cudart().cudaProfilerStart() output = block(block_inputs[index], rotation, *adaln_values, segments) torch.cuda.synchronize() torch.cuda.cudart().cudaProfilerStop() checksums[str(index)] = float(output.float().sum()) result = { "device": torch.cuda.get_device_name(), "torch": torch.__version__, "workload": workload, "profile_mode": args.profile_mode, "blocks": args.blocks, "checksums": checksums, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(result, indent=2) + "\n") print(json.dumps(result, indent=2), flush=True) return for index in args.blocks: hidden = block_inputs[index] expected = block_outputs[index] block = model.backbone.blocks[index] adaln_values = model.backbone.adaln[index](timesteps) def run(fused: bool): set_fused(index, fused) return block(hidden, rotation, *adaln_values, segments) for _ in range(args.warmup): run(False) run(True) torch.cuda.synchronize() baseline = run(False) candidate = run(True) torch.cuda.synchronize() parity = compare(candidate, baseline) baseline_capture = compare(baseline, expected) candidate_capture = compare(candidate, expected) del baseline, candidate samples = {"baseline": [], "candidate": []} for round_index in range(args.rounds): order = (False, True) if round_index % 2 == 0 else (True, False) for fused in order: elapsed, output = timed(lambda fused=fused: run(fused)) samples["candidate" if fused else "baseline"].append(elapsed) del output os.environ["H3_SAGE_ENTRY_FUSION"] = "0" baseline_memory = memory(lambda: run(False)) os.environ["H3_SAGE_ENTRY_FUSION"] = "1" candidate_memory = memory(lambda: run(True)) os.environ["H3_SAGE_ENTRY_FUSION"] = "0" baseline_events, baseline_launches = profile_module_kernels(lambda: run(False), warmup=1, iterations=1, row_limit=1000) os.environ["H3_SAGE_ENTRY_FUSION"] = "1" candidate_events, candidate_launches = profile_module_kernels(lambda: run(True), warmup=1, iterations=1, row_limit=1000) baseline_launches["device_kernel_launches"] = sum( row["count"] for row in baseline_events if row["device_type"] == "DeviceType.CUDA" ) candidate_launches["device_kernel_launches"] = sum( row["count"] for row in candidate_events if row["device_type"] == "DeviceType.CUDA" ) baseline_timing = summarize(samples["baseline"]) candidate_timing = summarize(samples["candidate"]) rows.append({ "block": index, "candidate_vs_baseline": parity, "baseline_vs_capture": baseline_capture, "candidate_vs_capture": candidate_capture, "baseline": { "timing": baseline_timing, "memory": baseline_memory, "launches": baseline_launches, }, "candidate": { "timing": candidate_timing, "memory": candidate_memory, "launches": candidate_launches, }, "p50_improvement_percent": (baseline_timing["p50_s"] - candidate_timing["p50_s"]) / baseline_timing["p50_s"] * 100.0, }) del adaln_values result = { "device": torch.cuda.get_device_name(), "torch": torch.__version__, "workload": workload, "rounds": args.rounds, "blocks": rows, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(result, indent=2) + "\n") print(json.dumps(result, indent=2), flush=True) if __name__ == "__main__": main()