144 lines
5.9 KiB
Python
144 lines
5.9 KiB
Python
"""Validate projection-strided Sage2 NHD execution on real H3 blocks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
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.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
|
|
from profile_h3_block import summarize
|
|
|
|
|
|
def timed(fn, warmup: int, iterations: int) -> dict[str, float]:
|
|
with torch.inference_mode():
|
|
for _ in range(warmup):
|
|
fn()
|
|
values = []
|
|
for _ in range(iterations):
|
|
torch.cuda.synchronize()
|
|
started = time.perf_counter()
|
|
fn()
|
|
torch.cuda.synchronize()
|
|
values.append(time.perf_counter() - started)
|
|
return summarize(values)
|
|
|
|
|
|
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("--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("--blocks", nargs="+", type=int, default=(0, 24, 49))
|
|
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)
|
|
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,
|
|
)
|
|
|
|
results = []
|
|
original_layout = os.environ.get("H3_SAGE_QKV_LAYOUT")
|
|
try:
|
|
for block_index in args.blocks:
|
|
if not 0 <= block_index < 50:
|
|
raise ValueError("block indices must be in [0, 49]")
|
|
block = H3DiTBlock.from_checkpoint(checkpoint, block_index, attention_backend="sage2").eval()
|
|
block.fused_elementwise = True
|
|
adaln = H3CurveAdaLN.from_checkpoint(checkpoint, f"blocks.{block_index}.adaln_proj").eval()
|
|
modulation = tuple(value.detach() for value in adaln(timesteps))
|
|
run = lambda: block(hidden.clone(), rotation, *modulation, segments)
|
|
|
|
with torch.inference_mode():
|
|
os.environ["H3_SAGE_QKV_LAYOUT"] = "hnd"
|
|
expected = run()
|
|
os.environ["H3_SAGE_QKV_LAYOUT"] = "strided_nhd"
|
|
actual = run()
|
|
delta = actual.float() - expected.float()
|
|
parity = {
|
|
"equal": torch.equal(actual, expected),
|
|
"max_abs": delta.abs().max().item(),
|
|
"mean_abs": delta.abs().mean().item(),
|
|
"reference_checksum": expected.float().sum().item(),
|
|
"candidate_checksum": actual.float().sum().item(),
|
|
}
|
|
if not parity["equal"]:
|
|
raise RuntimeError(f"block {block_index} Sage NHD output is not bit-exact: {parity}")
|
|
|
|
os.environ["H3_SAGE_QKV_LAYOUT"] = "hnd"
|
|
hnd_timing = timed(run, args.warmup, args.iterations)
|
|
os.environ["H3_SAGE_QKV_LAYOUT"] = "strided_nhd"
|
|
nhd_timing = timed(run, args.warmup, args.iterations)
|
|
results.append({
|
|
"block": block_index,
|
|
"parity": parity,
|
|
"hnd_timing": hnd_timing,
|
|
"strided_nhd_timing": nhd_timing,
|
|
"p50_speedup": hnd_timing["p50_s"] / nhd_timing["p50_s"],
|
|
"p50_latency_reduction": 1.0 - nhd_timing["p50_s"] / hnd_timing["p50_s"],
|
|
})
|
|
print(
|
|
f"block {block_index}: exact, hnd={hnd_timing['p50_s'] * 1000:.3f}ms, "
|
|
f"strided_nhd={nhd_timing['p50_s'] * 1000:.3f}ms, speedup={results[-1]['p50_speedup']:.3f}x",
|
|
flush=True,
|
|
)
|
|
finally:
|
|
if original_layout is None:
|
|
os.environ.pop("H3_SAGE_QKV_LAYOUT", None)
|
|
else:
|
|
os.environ["H3_SAGE_QKV_LAYOUT"] = original_layout
|
|
|
|
report = {
|
|
"device": torch.cuda.get_device_name(),
|
|
"torch": torch.__version__,
|
|
"resolution": [args.width, args.height],
|
|
"frames": aligned_frames,
|
|
"packed_tokens": hidden.shape[0],
|
|
"steps": args.steps,
|
|
"sampler_step": args.sampler_step,
|
|
"seed": args.seed,
|
|
"attention": "sage2",
|
|
"fused_elementwise": True,
|
|
"warmup": args.warmup,
|
|
"iterations": args.iterations,
|
|
"results": results,
|
|
}
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
args.output.write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|