"""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.nvfp4_quant import nvfp4_activation_scale, vortex_native_quantize_nvfp4, vortex_quantize_nvfp4 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) 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 quantize_activation(flat_x: torch.Tensor, quantizer: str, scale: torch.Tensor | None, timings: dict[str, list[float]] | None = None): from comfy_kitchen.tensor import QuantizedTensor if quantizer == "comfy": return QuantizedTensor.from_float(flat_x, "TensorCoreNVFP4Layout") if quantizer == "vortex_recalculate": return vortex_quantize_nvfp4(flat_x, timings=timings) if quantizer == "vortex_precomputed_scale": if scale is None: raise ValueError("vortex_precomputed_scale requires a precomputed scale") return vortex_quantize_nvfp4(flat_x, scale=scale, timings=timings) if quantizer == "vortex_native": return vortex_native_quantize_nvfp4(flat_x, timings=timings) raise ValueError(f"Unsupported quantizer: {quantizer}") def profile_linear_stages(module: Nvfp4Linear, x: torch.Tensor, *, iterations: int, quantizer: str) -> dict[str, dict[str, float]]: stats: dict[str, list[float]] = {} precomputed_scale = None if quantizer == "vortex_precomputed_scale": precomputed_scale = timed(stats, "precomputed_scale_calibration", lambda: nvfp4_activation_scale(x.reshape(-1, module.in_features).contiguous())) 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: quantize_activation(flat_x, quantizer, precomputed_scale, stats)) 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 run_linear_with_quantizer(module: Nvfp4Linear, x: torch.Tensor, quantizer: str, precomputed_scale: torch.Tensor | None = None) -> torch.Tensor: original_shape = x.shape[:-1] flat_x = x.reshape(-1, module.in_features).contiguous() if module.pre_quant_scale is not None: flat_x = flat_x * module.pre_quant_scale.to(flat_x) packed_weight = module._packed_weight() bias = module.bias.to(flat_x) if module.bias is not None else None if module.full_precision_matrix_mult: output = functional.linear(flat_x, packed_weight.dequantize().to(flat_x), bias) else: output = functional.linear(quantize_activation(flat_x, quantizer, precomputed_scale), packed_weight, bias) return output[:flat_x.shape[0], :module.out_features].reshape(*original_shape, module.out_features) def profile_cuda_kernels(module: Nvfp4Linear, x: torch.Tensor, *, warmup: int, iterations: int, row_limit: int, quantizer: str) -> list[dict]: with torch.inference_mode(): precomputed_scale = nvfp4_activation_scale(x.reshape(-1, module.in_features).contiguous()) if quantizer == "vortex_precomputed_scale" else None for _ in range(warmup): run_linear_with_quantizer(module, x, quantizer, precomputed_scale) 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): run_linear_with_quantizer(module, x, quantizer, precomputed_scale) 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("--quantizers", nargs="+", choices=("comfy", "vortex_recalculate", "vortex_precomputed_scale", "vortex_native"), default=("comfy", "vortex_recalculate", "vortex_precomputed_scale")) 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] reference = run_linear_with_quantizer(module, x, "comfy") for quantizer in args.quantizers: precomputed_scale = nvfp4_activation_scale(x.reshape(-1, module.in_features).contiguous()) if quantizer == "vortex_precomputed_scale" else None for _ in range(args.warmup): run_linear_with_quantizer(module, x, quantizer, precomputed_scale) stage_timings = profile_linear_stages(module, x, iterations=args.iterations, quantizer=quantizer) candidate = run_linear_with_quantizer(module, x, quantizer, precomputed_scale) diff = (candidate.float() - reference.float()).abs() kernels = profile_cuda_kernels(module, x, warmup=args.warmup, iterations=args.profiler_iterations, row_limit=args.profiler_row_limit, quantizer=quantizer) results.append( { "name": name, "quantizer": quantizer, "module": module_info(module, x), "stage_timings": stage_timings, "reference_diff": {"max": diff.max().item(), "mean": diff.mean().item()}, "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, "quantizers": args.quantizers, "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()