From d6eabb150ea7a0acf0e123e9fe4ace14db06e315 Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Sat, 15 Aug 2026 01:37:42 +0700 Subject: [PATCH] Add NVFP4 linear profiler --- tools/profile_nvfp4_linear.py | 248 ++++++++++++++++++++++++++++++++++ 1 file changed, 248 insertions(+) create mode 100644 tools/profile_nvfp4_linear.py diff --git a/tools/profile_nvfp4_linear.py b/tools/profile_nvfp4_linear.py new file mode 100644 index 0000000..230964e --- /dev/null +++ b/tools/profile_nvfp4_linear.py @@ -0,0 +1,248 @@ +"""Profile current NVFP4 linear execution stages and CUDA kernels.""" + +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 torch.profiler import ProfilerActivity, profile + +from h3_blackwell_runtime.adaln import H3CurveAdaLN +from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, rms_norm, rms_rope_split_half_, run_attention +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.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 + + +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_linear_stages(module: Nvfp4Linear, x: torch.Tensor, *, iterations: int) -> dict[str, dict[str, float]]: + from comfy_kitchen.tensor import QuantizedTensor + + stats: dict[str, list[float]] = {} + for _ in range(iterations): + original_shape = x.shape[:-1] + flat_x = timed(stats, "flatten_contiguous", lambda: x.reshape(-1, module.in_features).contiguous()) + if module.pre_quant_scale is not None: + flat_x = timed(stats, "pre_quant_scale", lambda: flat_x * module.pre_quant_scale.to(flat_x)) + else: + stats.setdefault("pre_quant_scale", []).append(0.0) + packed_weight = timed(stats, "packed_weight_wrapper", module._packed_weight) + bias = timed(stats, "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, "weight_dequantize", lambda: packed_weight.dequantize().to(flat_x)) + output = timed(stats, "linear", lambda: functional.linear(flat_x, weight, bias)) + else: + packed_x = timed(stats, "activation_quantize", lambda: QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout")) + output = timed(stats, "linear", lambda: functional.linear(packed_x, packed_weight, bias)) + timed(stats, "slice_reshape", lambda: output[:flat_x.shape[0], :module.out_features].reshape(*original_shape, module.out_features)) + return {name: summarize(values) for name, values in stats.items()} + + +def profile_cuda_kernels(module: Nvfp4Linear, x: torch.Tensor, *, warmup: int, iterations: int, row_limit: int) -> list[dict]: + with torch.inference_mode(): + for _ in range(warmup): + module(x) + sync() + activities = [ProfilerActivity.CPU] + if torch.cuda.is_available(): + activities.append(ProfilerActivity.CUDA) + with profile(activities=activities, record_shapes=True) as prof: + for _ in range(iterations): + module(x) + sync() + rows = [] + for event in prof.key_averages(group_by_input_shape=True): + cpu_us = float(getattr(event, "cpu_time_total", 0.0) or 0.0) + cuda_us = float(getattr(event, "cuda_time_total", 0.0) or 0.0) + rows.append( + { + "key": event.key, + "count": int(event.count), + "cpu_time_total_us": cpu_us, + "cuda_time_total_us": cuda_us, + "input_shapes": str(getattr(event, "input_shapes", "")), + } + ) + rows.sort(key=lambda item: (item["cuda_time_total_us"], item["cpu_time_total_us"]), reverse=True) + return rows[:row_limit] + + +def representative_inputs(args: argparse.Namespace) -> tuple[H3DiTBlock, dict[str, torch.Tensor], dict]: + 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) + sigma = beta_sigmas(args.steps, device=args.device)[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) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, _gate_mlp = adaln(timesteps) + + with torch.inference_mode(): + h_msa = modulate_segments(rms_norm(hidden, block.norm1_weight, block.norm_eps), shift_msa, scale_msa, segments) + sequence = h_msa.shape[0] + inner = block.attention.heads * block.attention.head_dim + qkv = block.attention.qkv_proj(h_msa) + q, k, v = qkv.split(inner, dim=-1) + q = q.view(1, sequence, block.attention.heads, block.attention.head_dim) + k = k.view(1, sequence, block.attention.heads, block.attention.head_dim) + v = v.view(1, sequence, block.attention.heads, block.attention.head_dim) + q, k = rms_rope_split_half_(q, k, rotation, block.attention.q_norm_weight, block.attention.k_norm_weight, block.attention.eps) + attn_out = run_attention(q.transpose(1, 2).contiguous(), k.transpose(1, 2).contiguous(), v.transpose(1, 2).contiguous(), backend=block.attention.backend, is_causal=False) + out_proj_input = attn_out.transpose(1, 2).reshape(sequence, inner).contiguous() + x_after_attn = gate_segments(hidden, block.attention.out_proj(out_proj_input), gate_msa, segments) + h_mlp = modulate_segments(rms_norm(x_after_attn, block.norm2_weight, block.norm_eps), shift_mlp, scale_mlp, segments) + gate, up = block.mlp.fc1(h_mlp).chunk(2, dim=-1) + fc2_input = torch.nn.functional.silu(gate).mul_(up) + + inputs = { + "attn_qkv_proj": h_msa, + "attn_out_proj": out_proj_input, + "mlp_fc1": h_mlp, + "mlp_fc2": fc2_input, + } + metadata = { + "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, + "hidden_shape": list(hidden.shape), + "segments": segments, + } + return block, inputs, metadata + + +def module_for_name(block: H3DiTBlock, name: str) -> Nvfp4Linear: + modules = { + "attn_qkv_proj": block.attention.qkv_proj, + "attn_out_proj": block.attention.out_proj, + "mlp_fc1": block.mlp.fc1, + "mlp_fc2": block.mlp.fc2, + } + return modules[name] + + +def module_info(module: Nvfp4Linear, x: torch.Tensor) -> dict: + return { + "class": type(module).__name__, + "in_features": module.in_features, + "out_features": module.out_features, + "output_dtype": str(module.output_dtype), + "full_precision_matrix_mult": module.full_precision_matrix_mult, + "weight_dtype": str(module.weight.dtype), + "weight_shape": list(module.weight.shape), + "weight_scale_dtype": str(module.weight_scale.dtype), + "weight_scale_shape": list(module.weight_scale.shape), + "weight_scale_2_dtype": str(module.weight_scale_2.dtype), + "weight_scale_2_shape": list(module.weight_scale_2.shape), + "bias_dtype": str(module.bias.dtype) if module.bias is not None else None, + "bias_shape": list(module.bias.shape) if module.bias is not None else None, + "pre_quant_scale": module.pre_quant_scale is not None, + "input_dtype": str(x.dtype), + "input_shape": list(x.shape), + "input_is_contiguous": x.is_contiguous(), + "input_stride": list(x.stride()), + } + + +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/nvfp4-linear-profile-7cc03f3.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) + parser.add_argument("--seed", type=int, default=440407) + parser.add_argument("--text-tokens", type=int, default=93) + parser.add_argument("--block-index", type=int, default=24) + parser.add_argument("--attention", choices=AVAILABLE_BACKENDS, default="sage2") + parser.add_argument("--linears", nargs="+", choices=("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2"), default=("attn_qkv_proj", "attn_out_proj", "mlp_fc1", "mlp_fc2")) + parser.add_argument("--warmup", type=int, default=2) + parser.add_argument("--iterations", type=int, default=5) + parser.add_argument("--profiler-iterations", type=int, default=2) + parser.add_argument("--profiler-row-limit", type=int, default=30) + parser.add_argument("--device", default="cuda") + return parser.parse_args() + + +def main() -> None: + args = parse_args() + block, inputs, metadata = representative_inputs(args) + results = [] + with torch.inference_mode(): + for name in args.linears: + module = module_for_name(block, name) + x = inputs[name] + for _ in range(args.warmup): + module(x) + stage_timings = profile_linear_stages(module, x, iterations=args.iterations) + kernels = profile_cuda_kernels(module, x, warmup=args.warmup, iterations=args.profiler_iterations, row_limit=args.profiler_row_limit) + results.append( + { + "name": name, + "module": module_info(module, x), + "stage_timings": stage_timings, + "profiler_top_events": kernels, + } + ) + + output = { + "model_path": args.model_path, + "block_index": args.block_index, + "attention": args.attention, + "warmup": args.warmup, + "iterations": args.iterations, + "profiler_iterations": args.profiler_iterations, + "metadata": metadata, + "results": results, + } + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(json.dumps(output, indent=2), encoding="utf-8") + print(json.dumps(output, indent=2), flush=True) + + +if __name__ == "__main__": + main()