h3-blackwell-runtime/tools/profile_h3_block.py

"""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()