"""Profile one representative H3 DiT block's dense/attention hot paths.""" from __future__ import annotations import argparse import json import time import warnings from pathlib import Path warnings.filterwarnings("ignore", message="Found GPU0 NVIDIA GB10 which is of cuda capability 12.1.*", category=UserWarning) import torch from h3_blackwell_runtime.adaln import H3CurveAdaLN from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, rms_rope_split_half_, rms_norm, run_attention from h3_blackwell_runtime.block import H3DiTBlock, gate_segments, modulate_segments from h3_blackwell_runtime.checkpoint import H3Checkpoint from h3_blackwell_runtime.packing import H3PromptPacker from h3_blackwell_runtime.rope import h3_rope_rotation from h3_blackwell_runtime.sampler import beta_sigmas, _audio_sigma, _model_sigma from h3_blackwell_runtime.t2v import random_av_latents CAT_PROMPT = ( "A playful orange tabby cat starts in an ordinary cozy living room in a normal house, afternoon light, sofa and rug. " "The cat crouches, jumps, and does one clean athletic backflip in slow motion. As the backflip completes there is a " "sharp cinematic cut: the cat lands perfectly on a glowing neon disco dance floor wearing oversized black sunglasses. " "Mirror ball reflections, colorful lights, joyful party energy, stylish and funny, clear before-and-after transformation." ) def sync() -> None: if torch.cuda.is_available(): torch.cuda.synchronize() def timed(stats: dict[str, list[float]], name: str, fn): sync() started = time.perf_counter() value = fn() sync() stats.setdefault(name, []).append(time.perf_counter() - started) return value def summarize(values: list[float]) -> dict[str, float]: ordered = sorted(values) return { "count": len(values), "mean_s": sum(values) / len(values), "min_s": ordered[0], "max_s": ordered[-1], } def profile_block( block: H3DiTBlock, hidden: torch.Tensor, rotation: torch.Tensor, adaln_values: tuple[torch.Tensor, ...], segments: list[tuple[int, int, int]], *, iterations: int, ) -> tuple[dict[str, dict[str, float]], torch.Tensor]: stats: dict[str, list[float]] = {} shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = adaln_values result = hidden for _ in range(iterations): x = hidden.clone() def run_once() -> torch.Tensor: nonlocal x h = timed(stats, "norm1", lambda: rms_norm(x, block.norm1_weight, block.norm_eps)) h = timed(stats, "modulate_msa", lambda: modulate_segments(h, shift_msa, scale_msa, segments)) attention = block.attention sequence = h.shape[0] inner = attention.heads * attention.head_dim qkv = timed(stats, "attn_qkv_proj", lambda: attention.qkv_proj(h)) q, k, v = timed(stats, "attn_qkv_split_view", lambda: tuple(t.view(1, sequence, attention.heads, attention.head_dim) for t in qkv.split(inner, dim=-1))) q, k = timed(stats, "attn_qk_rms_rope", lambda: rms_rope_split_half_(q, k, rotation, attention.q_norm_weight, attention.k_norm_weight, attention.eps)) q = timed(stats, "attn_q_transpose_contiguous", lambda: q.transpose(1, 2).contiguous()) k = timed(stats, "attn_k_transpose_contiguous", lambda: k.transpose(1, 2).contiguous()) v = timed(stats, "attn_v_transpose_contiguous", lambda: v.transpose(1, 2).contiguous()) attn_out = timed(stats, "attention_kernel", lambda: run_attention(q, k, v, backend=attention.backend, is_causal=False)) attn_rows = timed(stats, "attn_output_reshape", lambda: attn_out.transpose(1, 2).reshape(sequence, inner).contiguous()) attn_update = timed(stats, "attn_out_proj", lambda: attention.out_proj(attn_rows)) x = timed(stats, "gate_msa", lambda: gate_segments(x, attn_update, gate_msa, segments)) h2 = timed(stats, "norm2", lambda: rms_norm(x, block.norm2_weight, block.norm_eps)) h2 = timed(stats, "modulate_mlp", lambda: modulate_segments(h2, shift_mlp, scale_mlp, segments)) gate, up = timed(stats, "mlp_fc1", lambda: block.mlp.fc1(h2).chunk(2, dim=-1)) activated = timed(stats, "mlp_swiglu", lambda: torch.nn.functional.silu(gate).mul_(up)) mlp_update = timed(stats, "mlp_fc2", lambda: block.mlp.fc2(activated)) x = timed(stats, "gate_mlp", lambda: gate_segments(x, mlp_update, gate_mlp, segments)) return x result = timed(stats, "block_total", run_once) return {name: summarize(values) for name, values in stats.items()}, result def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors") parser.add_argument("--output", type=Path, default=Path("/output/h3-blackwell-runtime/benchmarks/h3-block-profile-3d0c093.json")) parser.add_argument("--width", type=int, default=960) parser.add_argument("--height", type=int, default=544) 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, help="One-indexed sampler step used to build representative timesteps.") parser.add_argument("--seed", type=int, default=440407) parser.add_argument("--text-tokens", type=int, default=93, help="Synthetic refined-text token count; avoids Qwen/refiner load.") parser.add_argument("--block-index", type=int, default=24) parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2") parser.add_argument("--warmup", type=int, default=1) parser.add_argument("--iterations", type=int, default=3) parser.add_argument("--device", default="cuda") return parser.parse_args() def main() -> None: args = parse_args() if not 0 <= args.block_index < 50: raise ValueError("--block-index must be in [0, 49]") if args.sampler_step < 1 or args.sampler_step > args.steps: raise ValueError("--sampler-step must be between 1 and --steps") torch.manual_seed(args.seed) checkpoint = H3Checkpoint(args.model_path, device=args.device) block = H3DiTBlock.from_checkpoint(checkpoint, args.block_index, attention_backend=args.attention).eval() adaln = H3CurveAdaLN.from_checkpoint(checkpoint, f"blocks.{args.block_index}.adaln_proj").eval() packer = H3PromptPacker(checkpoint) video, audio, aligned_frames = random_av_latents(args.width, args.height, args.frames, args.seed, device=args.device) sigmas = beta_sigmas(args.steps, device=args.device) sigma = sigmas[args.sampler_step - 1] native_audio = audio.to(torch.bfloat16) * (_audio_sigma(sigma) / sigma) text = torch.randn(1, args.text_tokens, 5376, device=args.device, dtype=torch.bfloat16) hidden, timesteps, segments, positions, _, _ = packer(text, video, native_audio, _model_sigma(sigma)) rotation = h3_rope_rotation(positions.to(args.device), checkpoint.tensor("rope.inv_freq", dtype=torch.float32), hidden.dtype) adaln_values = tuple(value.detach() for value in adaln(timesteps)) with torch.inference_mode(): for _ in range(args.warmup): profile_block(block, hidden, rotation, adaln_values, segments, iterations=1) stats, output = profile_block(block, hidden, rotation, adaln_values, segments, iterations=args.iterations) result = { "prompt": CAT_PROMPT, "model_path": args.model_path, "width": args.width, "height": args.height, "frames": aligned_frames, "steps": args.steps, "sampler_step": args.sampler_step, "seed": args.seed, "text_tokens": args.text_tokens, "block_index": args.block_index, "attention": args.attention, "warmup": args.warmup, "iterations": args.iterations, "hidden_shape": list(hidden.shape), "output_shape": list(output.shape), "segments": segments, "timings": stats, } args.output.parent.mkdir(parents=True, exist_ok=True) args.output.write_text(json.dumps(result, indent=2), encoding="utf-8") print(json.dumps(result, indent=2), flush=True) if __name__ == "__main__": main()