"""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 import torch.nn.functional as functional from h3_blackwell_runtime.adaln import H3CurveAdaLN from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, rms_rope_split_half_, rms_norm, run_attention, run_sol_attention_bshd from h3_blackwell_runtime.block import H3DiTBlock, gate_segments, modulate_segments from h3_blackwell_runtime.checkpoint import H3Checkpoint from h3_blackwell_runtime.nvfp4 import Nvfp4Linear from h3_blackwell_runtime.nvfp4_quant import nvfp4_activation_scale 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) def percentile(percent: float) -> float: if len(ordered) == 1: return ordered[0] rank = (len(ordered) - 1) * percent low = int(rank) high = min(low + 1, len(ordered) - 1) weight = rank - low return ordered[low] * (1.0 - weight) + ordered[high] * weight return { "count": len(values), "mean_s": sum(values) / len(values), "p50_s": percentile(0.50), "p90_s": percentile(0.90), "p95_s": percentile(0.95), "p99_s": percentile(0.99), "min_s": ordered[0], "max_s": ordered[-1], } def profiled_nvfp4_linear(stats: dict[str, list[float]], prefix: str, module: Nvfp4Linear, x: torch.Tensor) -> torch.Tensor: from comfy_kitchen.tensor import QuantizedTensor, TensorCoreNVFP4Layout original_shape = x.shape[:-1] flat_x = timed(stats, f"{prefix}.flatten_contiguous", lambda: x.reshape(-1, module.in_features).contiguous()) if module.pre_quant_scale is not None: flat_x = timed(stats, f"{prefix}.pre_quant_scale", lambda: flat_x * module.pre_quant_scale.to(flat_x)) else: stats.setdefault(f"{prefix}.pre_quant_scale", []).append(0.0) packed_weight = timed(stats, f"{prefix}.packed_weight_wrapper", module._packed_weight) bias = timed(stats, f"{prefix}.bias_cast", lambda: module.bias.to(flat_x) if module.bias is not None else None) if module.full_precision_matrix_mult: weight = timed(stats, f"{prefix}.weight_dequantize", lambda: packed_weight.dequantize().to(flat_x)) output = timed(stats, f"{prefix}.gemm", lambda: functional.linear(flat_x, weight, bias)) return timed(stats, f"{prefix}.slice_reshape", lambda: output.reshape(*original_shape, module.out_features)) if flat_x.dtype == torch.float32: raise ValueError("Quantized NVFP4 activation GEMM requires FP16 or BF16 activations.") orig_shape = tuple(flat_x.shape) scale = timed(stats, f"{prefix}.activation_scale", lambda: nvfp4_activation_scale(flat_x)) scale = timed(stats, f"{prefix}.scale_to_device", lambda: torch.as_tensor(scale, device=flat_x.device, dtype=torch.float32)) qdata, block_scale = timed( stats, f"{prefix}.activation_quant_pack", lambda: __import__("comfy_kitchen").quantize_nvfp4( flat_x, scale, pad_16x=TensorCoreNVFP4Layout.get_padded_shape(orig_shape) != orig_shape, ), ) packed_x = timed( stats, f"{prefix}.activation_quant_wrap", lambda: QuantizedTensor( qdata, "TensorCoreNVFP4Layout", TensorCoreNVFP4Layout.Params( scale=scale, orig_dtype=flat_x.dtype, orig_shape=orig_shape, block_scale=block_scale, ), ), ) output = timed(stats, f"{prefix}.gemm", lambda: functional.linear(packed_x, packed_weight, bias)) return timed(stats, f"{prefix}.slice_reshape", lambda: output[:flat_x.shape[0], :module.out_features].reshape(*original_shape, module.out_features)) 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: profiled_nvfp4_linear(stats, "linear.attn_qkv_proj", 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)) if attention.backend == "sol_attn": q = timed(stats, "attn_q_bshd_contiguous", lambda: q.contiguous()) k = timed(stats, "attn_k_bshd_contiguous", lambda: k.contiguous()) v = timed(stats, "attn_v_bshd_contiguous", lambda: v.contiguous()) attn_out = timed(stats, "attention_kernel", lambda: run_sol_attention_bshd(q, k, v, is_causal=False)) attn_rows = timed(stats, "attn_output_reshape", lambda: attn_out.reshape(sequence, inner).contiguous()) else: 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: profiled_nvfp4_linear(stats, "linear.attn_out_proj", 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: profiled_nvfp4_linear(stats, "linear.mlp_fc1", 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: profiled_nvfp4_linear(stats, "linear.mlp_fc2", 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="sol_attn") parser.add_argument("--warmup", type=int, default=20) parser.add_argument("--iterations", type=int, default=50) 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()