h3-blackwell-runtime/tools/profile_nvfp4_linear.py

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