135 lines
5.3 KiB
Python
135 lines
5.3 KiB
Python
"""Validate fused SwiGLU and native NVFP4 production on real H3 FC1 output."""
|
|
|
|
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.block import H3DiTBlock
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
|
from h3_blackwell_runtime.nvfp4_quant import (
|
|
nvfp4_activation_scale,
|
|
vortex_native_quantize_swiglu_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 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, required=True)
|
|
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, 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,
|
|
)
|
|
modulation = tuple(value.detach() for value in adaln(timesteps))
|
|
captured = []
|
|
block.fused_nvfp4_modulation = False
|
|
hook = block.mlp.fc1.register_forward_hook(
|
|
lambda _module, _inputs, output: captured.append(output.detach()),
|
|
)
|
|
with torch.inference_mode():
|
|
block(hidden, rotation, *modulation, segments)
|
|
hook.remove()
|
|
gate_up = captured[-1]
|
|
|
|
with torch.inference_mode():
|
|
gate, up = gate_up.chunk(2, dim=-1)
|
|
materialized = torch.nn.functional.silu(gate).mul_(up)
|
|
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),
|
|
)
|
|
actual_scale, actual_qdata, actual_sfa = vortex_native_quantize_swiglu_nvfp4(
|
|
gate_up,
|
|
)
|
|
torch.cuda.synchronize()
|
|
for _ in range(args.warmup):
|
|
vortex_native_quantize_swiglu_nvfp4(gate_up)
|
|
times = []
|
|
for _ in range(args.iterations):
|
|
torch.cuda.synchronize()
|
|
started = time.perf_counter()
|
|
vortex_native_quantize_swiglu_nvfp4(gate_up)
|
|
torch.cuda.synchronize()
|
|
times.append(time.perf_counter() - started)
|
|
|
|
report = {
|
|
"device": torch.cuda.get_device_name(),
|
|
"block_index": args.block_index,
|
|
"input_shape": list(gate_up.shape),
|
|
"output_shape": list(materialized.shape),
|
|
"scale_equal": torch.equal(actual_scale, reference_scale),
|
|
"scale_reference": reference_scale.item(),
|
|
"scale_actual": actual_scale.item(),
|
|
"qdata_differences": torch.count_nonzero(actual_qdata != reference_qdata).item(),
|
|
"sfa_differences": torch.count_nonzero(
|
|
actual_sfa.view(torch.uint8) != reference_sfa.view(torch.uint8)
|
|
).item(),
|
|
"producer_p50_ms": sorted(times)[len(times) // 2] * 1000.0,
|
|
}
|
|
report["equal"] = (
|
|
report["scale_equal"]
|
|
and report["qdata_differences"] == 0
|
|
and report["sfa_differences"] == 0
|
|
)
|
|
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 SwiGLU producer is not byte-exact")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|