diff --git a/PLAN.md b/PLAN.md index 94893bc..c1c340c 100644 --- a/PLAN.md +++ b/PLAN.md @@ -63,6 +63,37 @@ Prompt-only FL2VA is now at warm Comfy parity with the direct Sage2 baseline. Fe - Initial direct cache modes are implemented as opt-in approximate sampler modes. They reuse cached denoised deltas and report skipped-step stats; full-size quality/threshold sweeps are still required before using them for production output. 6. Keep every backend explicit per run, with separate quality and timing records for sampling, VAE, audio, and end-to-end output. +## Next Performance Plan: Blackwell NVFP4 GEMMs + +The latest attention/cache sweeps show that attention backend swaps are not the main remaining speed lever. Full request times for `sage2`, `kj_sage_fp8`, `kj_sage_fp8pp`, `kj_head_sliced`, and `sol_attn` are close, while cache gains come from skipping denoiser calls and must remain quality-gated. The next exact performance target is therefore the dense transformer linear stack. + +Target mixed-precision policy: + +- Keep residual stream, RMSNorm, RoPE, modulation/AdaLN, residual adds, timestep embeddings, softmax, final layer, audio VAE, and video VAE in BF16/FP16. +- Keep attention accumulation/output in the best measured BF16/FP16 backend until a separate quality gate proves otherwise. +- Replace only the large transformer GEMMs first: QKV projection, attention output projection, MLP gate/up projection, and MLP down projection. +- Prioritize Blackwell-native NVFP4 activation x NVFP4 weight GEMMs with BF16 accumulation/output, avoiding unnecessary dequantize/contiguous boundaries. + +Kernel candidates to test, in order: + +1. Current `Nvfp4Linear`/Comfy Kitchen path as the measured baseline. +2. Transformer Engine NVFP4BlockScaling prototype to establish expected Blackwell FP4 behavior with a higher-level NVIDIA stack. +3. CUTLASS/CuTe block-scaled NVFP4 GEMM prototype for owned hot-path kernels, exact layouts, fused epilogues, and persistent-kernel experiments. +4. cuBLASLt NVFP4/grouped GEMM where CUDA exposes a supported operation on the current Blackwell target. + +Validation and profiling sequence: + +1. Use `tools/profile_h3_block.py` to measure one representative H3 block before writing kernels. Record QKV, RoPE/RMS, attention kernel, output projection, MLP fc1, activation, MLP fc2, modulation/gating, and total block time. +2. Repeat profiling for representative block indices, at least block `0`, `24`, and `49`, because token statistics and cache behavior can differ through depth. +3. Confirm which component dominates before implementing a kernel replacement. +4. Add candidate GEMM backends behind the existing `Nvfp4Linear` API so model code and correctness tests remain stable. +5. Gate each candidate by layer-level max/mean error, one-block output error, one denoiser-step tensor error, and finally full video quality. + +DGX Spark caveat: + +- GB10 reports SM121, not B200/GB200 SM100. CUTLASS, Transformer Engine, and cuBLASLt FP4 coverage must be probed independently on Spark before assuming B200 examples work unchanged. +- Keep separate kernel policy notes for GB10/Spark and B200/GB200-class Blackwell. + ## Non-Negotiable Validation - Never silently pad semantic H3 tokens for unmasked attention. diff --git a/tools/profile_h3_block.py b/tools/profile_h3_block.py new file mode 100644 index 0000000..bf9dce0 --- /dev/null +++ b/tools/profile_h3_block.py @@ -0,0 +1,176 @@ +"""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 + +from h3_blackwell_runtime.adaln import H3CurveAdaLN +from h3_blackwell_runtime.attention import AVAILABLE_BACKENDS, rms_rope_split_half_, rms_norm, run_attention +from h3_blackwell_runtime.block import H3DiTBlock, gate_segments, modulate_segments +from h3_blackwell_runtime.checkpoint import H3Checkpoint +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) + return { + "count": len(values), + "mean_s": sum(values) / len(values), + "min_s": ordered[0], + "max_s": ordered[-1], + } + + +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: 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)) + 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: 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: 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: 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="sage2") + parser.add_argument("--warmup", type=int, default=1) + parser.add_argument("--iterations", type=int, default=3) + 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()