146 lines
6 KiB
Diff
146 lines
6 KiB
Diff
diff --git a/tools/validate_nvfp4_modulate_producer.py b/tools/validate_nvfp4_modulate_producer.py
|
|
new file mode 100644
|
|
index 0000000..a17b7f2
|
|
--- /dev/null
|
|
+++ b/tools/validate_nvfp4_modulate_producer.py
|
|
@@ -0,0 +1,140 @@
|
|
+"""Validate fused H3 modulation and native NVFP4 production."""
|
|
+
|
|
+from __future__ import annotations
|
|
+
|
|
+import argparse
|
|
+import json
|
|
+import time
|
|
+from pathlib import Path
|
|
+
|
|
+import torch
|
|
+
|
|
+from h3_blackwell_runtime.adaln import H3CurveAdaLN
|
|
+from h3_blackwell_runtime.attention import rms_norm
|
|
+from h3_blackwell_runtime.block import H3DiTBlock, modulate_segments
|
|
+from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
+from h3_blackwell_runtime.h3_fusion import segment_index
|
|
+from h3_blackwell_runtime.nvfp4_quant import (
|
|
+ nvfp4_activation_scale,
|
|
+ vortex_native_quantize_modulated_nvfp4,
|
|
+)
|
|
+from h3_blackwell_runtime.packing import H3PromptPacker
|
|
+from h3_blackwell_runtime.sampler import _audio_sigma, _model_sigma, beta_sigmas
|
|
+from h3_blackwell_runtime.t2v import random_av_latents
|
|
+
|
|
+
|
|
+def parse_args() -> argparse.Namespace:
|
|
+ parser = argparse.ArgumentParser(description=__doc__)
|
|
+ parser.add_argument("--model-path", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
|
+ parser.add_argument("--output", type=Path, required=True)
|
|
+ parser.add_argument("--block-index", type=int, default=24)
|
|
+ parser.add_argument("--rows", type=int, default=2048)
|
|
+ parser.add_argument("--width", type=int, default=1344)
|
|
+ parser.add_argument("--height", type=int, default=768)
|
|
+ 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=440420)
|
|
+ parser.add_argument("--text-tokens", type=int, default=100)
|
|
+ parser.add_argument("--warmup", type=int, default=3)
|
|
+ parser.add_argument("--iterations", type=int, default=10)
|
|
+ parser.add_argument("--device", default="cuda")
|
|
+ return parser.parse_args()
|
|
+
|
|
+
|
|
+def main() -> None:
|
|
+ args = parse_args()
|
|
+ torch.manual_seed(args.seed)
|
|
+ checkpoint = H3Checkpoint(args.model_path, device=args.device)
|
|
+ block = H3DiTBlock.from_checkpoint(checkpoint, args.block_index, attention_backend="sage2").eval()
|
|
+ adaln = H3CurveAdaLN.from_checkpoint(
|
|
+ checkpoint, f"blocks.{args.block_index}.adaln_proj",
|
|
+ ).eval()
|
|
+ packer = H3PromptPacker(checkpoint)
|
|
+ video, audio, _ = 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, _, _, _ = packer(
|
|
+ text, video, native_audio, _model_sigma(sigma),
|
|
+ )
|
|
+ rows = min(args.rows, hidden.shape[0])
|
|
+ hidden = hidden[:rows].contiguous()
|
|
+ clipped_segments = []
|
|
+ for start, stop, table_row in segments:
|
|
+ if start >= rows:
|
|
+ break
|
|
+ clipped_segments.append((start, min(stop, rows), table_row))
|
|
+ shift_msa, scale_msa, *_ = (value.detach() for value in adaln(timesteps))
|
|
+
|
|
+ with torch.inference_mode():
|
|
+ normalized = rms_norm(hidden, block.norm1_weight, block.norm_eps)
|
|
+ materialized = modulate_segments(
|
|
+ normalized, shift_msa, scale_msa, clipped_segments,
|
|
+ )
|
|
+ reference_scale = nvfp4_activation_scale(materialized).float()
|
|
+ import comfy_kitchen as ck
|
|
+ from comfy_kitchen.tensor import TensorCoreNVFP4Layout
|
|
+
|
|
+ reference_qdata, reference_sfa = ck.quantize_nvfp4(
|
|
+ materialized,
|
|
+ reference_scale,
|
|
+ pad_16x=TensorCoreNVFP4Layout.get_padded_shape(tuple(materialized.shape))
|
|
+ != tuple(materialized.shape),
|
|
+ )
|
|
+ fused_scale, fused_qdata, fused_sfa = vortex_native_quantize_modulated_nvfp4(
|
|
+ normalized,
|
|
+ shift_msa.contiguous(),
|
|
+ scale_msa.contiguous(),
|
|
+ segment_index(rows, clipped_segments, hidden.device),
|
|
+ )
|
|
+ torch.cuda.synchronize()
|
|
+
|
|
+ def run_fused():
|
|
+ return vortex_native_quantize_modulated_nvfp4(
|
|
+ normalized,
|
|
+ shift_msa.contiguous(),
|
|
+ scale_msa.contiguous(),
|
|
+ segment_index(rows, clipped_segments, hidden.device),
|
|
+ )
|
|
+
|
|
+ for _ in range(args.warmup):
|
|
+ run_fused()
|
|
+ fused_times = []
|
|
+ for _ in range(args.iterations):
|
|
+ torch.cuda.synchronize()
|
|
+ started = time.perf_counter()
|
|
+ run_fused()
|
|
+ torch.cuda.synchronize()
|
|
+ fused_times.append(time.perf_counter() - started)
|
|
+
|
|
+ report = {
|
|
+ "device": torch.cuda.get_device_name(),
|
|
+ "block_index": args.block_index,
|
|
+ "rows": rows,
|
|
+ "width": hidden.shape[1],
|
|
+ "scale_equal": torch.equal(fused_scale, reference_scale),
|
|
+ "scale_reference": reference_scale.item(),
|
|
+ "scale_fused": fused_scale.item(),
|
|
+ "qdata_equal": torch.equal(fused_qdata, reference_qdata),
|
|
+ "qdata_differences": torch.count_nonzero(fused_qdata != reference_qdata).item(),
|
|
+ "sfa_equal": torch.equal(fused_sfa.view(torch.uint8), reference_sfa.view(torch.uint8)),
|
|
+ "sfa_differences": torch.count_nonzero(
|
|
+ fused_sfa.view(torch.uint8) != reference_sfa.view(torch.uint8)
|
|
+ ).item(),
|
|
+ "fused_producer_p50_ms": sorted(fused_times)[len(fused_times) // 2] * 1000.0,
|
|
+ }
|
|
+ report["equal"] = report["scale_equal"] and report["qdata_equal"] and report["sfa_equal"]
|
|
+ args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
+ args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
+ print(json.dumps(report, indent=2), flush=True)
|
|
+ if not report["equal"]:
|
|
+ raise RuntimeError("fused modulation producer is not byte-exact")
|
|
+
|
|
+
|
|
+if __name__ == "__main__":
|
|
+ main()
|