247 lines
12 KiB
Python
247 lines
12 KiB
Python
"""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()
|