290 lines
12 KiB
Python
290 lines
12 KiB
Python
|
"""Gate the fused Sage2 entry path on captured complete H3 blocks."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
import types
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
from h3_blackwell_runtime.checkpoint import H3Checkpoint
|
||
|
|
from h3_blackwell_runtime.denoiser import H3PackedDenoiser
|
||
|
|
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 profile_module_kernels, summarize
|
||
|
|
|
||
|
|
|
||
|
|
def compare(actual: torch.Tensor, expected: torch.Tensor) -> dict[str, float | int | bool]:
|
||
|
|
different = torch.count_nonzero(actual != expected).item()
|
||
|
|
delta = (actual.float() - expected.float()).abs()
|
||
|
|
return {
|
||
|
|
"equal": different == 0,
|
||
|
|
"different_elements": different,
|
||
|
|
"max_abs": float(delta.max()),
|
||
|
|
"mean_abs": float(delta.mean()),
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def timed(fn) -> tuple[float, torch.Tensor]:
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
started = time.perf_counter()
|
||
|
|
output = fn()
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
return time.perf_counter() - started, output
|
||
|
|
|
||
|
|
|
||
|
|
def memory(fn) -> dict[str, int]:
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
torch.cuda.reset_peak_memory_stats()
|
||
|
|
allocated = torch.cuda.memory_allocated()
|
||
|
|
reserved = torch.cuda.memory_reserved()
|
||
|
|
output = fn()
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
result = {
|
||
|
|
"allocated_bytes_before": allocated,
|
||
|
|
"reserved_bytes_before": reserved,
|
||
|
|
"peak_allocated_bytes": torch.cuda.max_memory_allocated(),
|
||
|
|
"peak_reserved_bytes": torch.cuda.max_memory_reserved(),
|
||
|
|
"peak_allocated_delta_bytes": torch.cuda.max_memory_allocated() - allocated,
|
||
|
|
"peak_reserved_delta_bytes": torch.cuda.max_memory_reserved() - reserved,
|
||
|
|
}
|
||
|
|
del output
|
||
|
|
return result
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser()
|
||
|
|
parser.add_argument("--capture", type=Path)
|
||
|
|
parser.add_argument("--model", default="/models/minimax_h3_fl2va_pruned_nvfp4.safetensors")
|
||
|
|
parser.add_argument("--blocks", type=int, nargs="+", default=[0, 24, 49])
|
||
|
|
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=2)
|
||
|
|
parser.add_argument("--rounds", type=int, default=8)
|
||
|
|
parser.add_argument("--cuda-profiler-capture", action="store_true")
|
||
|
|
parser.add_argument("--profile-mode", choices=("baseline", "candidate"), default="baseline")
|
||
|
|
parser.add_argument(
|
||
|
|
"--output",
|
||
|
|
type=Path,
|
||
|
|
default=Path("/output/h3-blackwell-runtime/benchmarks/gb10-sage2-p1-block-gate.json"),
|
||
|
|
)
|
||
|
|
args = parser.parse_args()
|
||
|
|
|
||
|
|
os.environ["H3_SAGE_QKV_LAYOUT"] = "strided_nhd"
|
||
|
|
os.environ["H3_SAGE_ENTRY_FUSION"] = "0"
|
||
|
|
torch.manual_seed(args.seed)
|
||
|
|
checkpoint = H3Checkpoint(args.model)
|
||
|
|
model = H3PackedDenoiser.from_checkpoint(checkpoint, attention_backend="sage2").eval()
|
||
|
|
if args.capture is not None:
|
||
|
|
captured = torch.load(args.capture / "input.pt", map_location="cuda", weights_only=False)
|
||
|
|
hidden = captured["hidden"].to("cuda")
|
||
|
|
timesteps = captured["timesteps"].to("cuda")
|
||
|
|
positions = captured["position_ids"].to("cuda")
|
||
|
|
segments = captured["segments"]
|
||
|
|
workload = {"capture": str(args.capture)}
|
||
|
|
else:
|
||
|
|
packer = H3PromptPacker(checkpoint)
|
||
|
|
video, audio, aligned_frames = random_av_latents(
|
||
|
|
args.width, args.height, args.frames, args.seed, device="cuda",
|
||
|
|
)
|
||
|
|
sigma = beta_sigmas(args.steps, device="cuda")[args.sampler_step - 1]
|
||
|
|
native_audio = audio.to(torch.bfloat16) * (_audio_sigma(sigma) / sigma)
|
||
|
|
text = torch.randn(1, args.text_tokens, 5376, device="cuda", dtype=torch.bfloat16)
|
||
|
|
hidden, timesteps, segments, positions, _, _ = packer(
|
||
|
|
text, video, native_audio, _model_sigma(sigma),
|
||
|
|
)
|
||
|
|
workload = {
|
||
|
|
"resolution": [args.width, args.height],
|
||
|
|
"frames": aligned_frames,
|
||
|
|
"steps": args.steps,
|
||
|
|
"sampler_step": args.sampler_step,
|
||
|
|
"seed": args.seed,
|
||
|
|
"text_tokens": args.text_tokens,
|
||
|
|
}
|
||
|
|
rotation = h3_rope_rotation(positions.to("cuda"), model.backbone.inv_freq, torch.bfloat16)
|
||
|
|
workload["tokens"] = hidden.shape[0]
|
||
|
|
workload["hidden_shape"] = list(hidden.shape)
|
||
|
|
workload["rotation_shape"] = list(rotation.shape)
|
||
|
|
print(workload, flush=True)
|
||
|
|
block_inputs = {}
|
||
|
|
block_outputs = {}
|
||
|
|
rows = []
|
||
|
|
|
||
|
|
with torch.inference_mode():
|
||
|
|
for index, (block, adaln) in enumerate(zip(model.backbone.blocks, model.backbone.adaln, strict=True)):
|
||
|
|
if index in args.blocks:
|
||
|
|
block_inputs[index] = hidden
|
||
|
|
hidden = block(hidden, rotation, *adaln(timesteps), segments)
|
||
|
|
if index in args.blocks:
|
||
|
|
block_outputs[index] = hidden
|
||
|
|
if index >= max(args.blocks):
|
||
|
|
break
|
||
|
|
|
||
|
|
baseline_forwards = {index: model.backbone.blocks[index].attention.forward for index in args.blocks}
|
||
|
|
|
||
|
|
def fused_attention_forward(
|
||
|
|
attention,
|
||
|
|
x,
|
||
|
|
rope_rotation,
|
||
|
|
sequence_parallel=None,
|
||
|
|
tensor_parallel=None,
|
||
|
|
modulation=None,
|
||
|
|
):
|
||
|
|
from h3_blackwell_runtime.sage2_entry import attention_nhd
|
||
|
|
|
||
|
|
if sequence_parallel is not None or tensor_parallel is not None:
|
||
|
|
raise ValueError("the Sage2 entry-fusion benchmark is single-GPU only")
|
||
|
|
sequence = x.shape[0]
|
||
|
|
inner = attention.heads * attention.head_dim
|
||
|
|
qkv = attention.qkv_proj.forward_modulated(x, *modulation) if modulation is not None else attention.qkv_proj(x)
|
||
|
|
q, k, v = qkv.split(inner, dim=-1)
|
||
|
|
q = q.view(1, sequence, attention.heads, attention.head_dim)
|
||
|
|
k = k.view(1, sequence, attention.heads, attention.head_dim)
|
||
|
|
v = v.view(1, sequence, attention.heads, attention.head_dim)
|
||
|
|
output = attention_nhd(
|
||
|
|
q,
|
||
|
|
k,
|
||
|
|
v,
|
||
|
|
rope_rotation,
|
||
|
|
attention.q_norm_weight,
|
||
|
|
attention.k_norm_weight,
|
||
|
|
attention.eps,
|
||
|
|
)
|
||
|
|
return attention.out_proj(output.reshape(sequence, inner))
|
||
|
|
|
||
|
|
candidate_forwards = {
|
||
|
|
index: types.MethodType(fused_attention_forward, model.backbone.blocks[index].attention)
|
||
|
|
for index in args.blocks
|
||
|
|
}
|
||
|
|
|
||
|
|
def set_fused(index: int, fused: bool) -> None:
|
||
|
|
model.backbone.blocks[index].attention.forward = (
|
||
|
|
candidate_forwards[index] if fused else baseline_forwards[index]
|
||
|
|
)
|
||
|
|
|
||
|
|
if args.cuda_profiler_capture:
|
||
|
|
checksums = {}
|
||
|
|
fused = args.profile_mode == "candidate"
|
||
|
|
os.environ["H3_SAGE_ENTRY_FUSION"] = "1" if fused else "0"
|
||
|
|
for index in args.blocks:
|
||
|
|
block = model.backbone.blocks[index]
|
||
|
|
set_fused(index, fused)
|
||
|
|
adaln_values = model.backbone.adaln[index](timesteps)
|
||
|
|
for _ in range(args.warmup):
|
||
|
|
block(block_inputs[index], rotation, *adaln_values, segments)
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
torch.cuda.cudart().cudaProfilerStart()
|
||
|
|
output = block(block_inputs[index], rotation, *adaln_values, segments)
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
torch.cuda.cudart().cudaProfilerStop()
|
||
|
|
checksums[str(index)] = float(output.float().sum())
|
||
|
|
result = {
|
||
|
|
"device": torch.cuda.get_device_name(),
|
||
|
|
"torch": torch.__version__,
|
||
|
|
"workload": workload,
|
||
|
|
"profile_mode": args.profile_mode,
|
||
|
|
"blocks": args.blocks,
|
||
|
|
"checksums": checksums,
|
||
|
|
}
|
||
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
args.output.write_text(json.dumps(result, indent=2) + "\n")
|
||
|
|
print(json.dumps(result, indent=2), flush=True)
|
||
|
|
return
|
||
|
|
|
||
|
|
for index in args.blocks:
|
||
|
|
hidden = block_inputs[index]
|
||
|
|
expected = block_outputs[index]
|
||
|
|
block = model.backbone.blocks[index]
|
||
|
|
adaln_values = model.backbone.adaln[index](timesteps)
|
||
|
|
|
||
|
|
def run(fused: bool):
|
||
|
|
set_fused(index, fused)
|
||
|
|
return block(hidden, rotation, *adaln_values, segments)
|
||
|
|
|
||
|
|
for _ in range(args.warmup):
|
||
|
|
run(False)
|
||
|
|
run(True)
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
|
||
|
|
baseline = run(False)
|
||
|
|
candidate = run(True)
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
parity = compare(candidate, baseline)
|
||
|
|
baseline_capture = compare(baseline, expected)
|
||
|
|
candidate_capture = compare(candidate, expected)
|
||
|
|
del baseline, candidate
|
||
|
|
|
||
|
|
samples = {"baseline": [], "candidate": []}
|
||
|
|
for round_index in range(args.rounds):
|
||
|
|
order = (False, True) if round_index % 2 == 0 else (True, False)
|
||
|
|
for fused in order:
|
||
|
|
elapsed, output = timed(lambda fused=fused: run(fused))
|
||
|
|
samples["candidate" if fused else "baseline"].append(elapsed)
|
||
|
|
del output
|
||
|
|
|
||
|
|
os.environ["H3_SAGE_ENTRY_FUSION"] = "0"
|
||
|
|
baseline_memory = memory(lambda: run(False))
|
||
|
|
os.environ["H3_SAGE_ENTRY_FUSION"] = "1"
|
||
|
|
candidate_memory = memory(lambda: run(True))
|
||
|
|
|
||
|
|
os.environ["H3_SAGE_ENTRY_FUSION"] = "0"
|
||
|
|
baseline_events, baseline_launches = profile_module_kernels(lambda: run(False), warmup=1, iterations=1, row_limit=1000)
|
||
|
|
os.environ["H3_SAGE_ENTRY_FUSION"] = "1"
|
||
|
|
candidate_events, candidate_launches = profile_module_kernels(lambda: run(True), warmup=1, iterations=1, row_limit=1000)
|
||
|
|
baseline_launches["device_kernel_launches"] = sum(
|
||
|
|
row["count"] for row in baseline_events if row["device_type"] == "DeviceType.CUDA"
|
||
|
|
)
|
||
|
|
candidate_launches["device_kernel_launches"] = sum(
|
||
|
|
row["count"] for row in candidate_events if row["device_type"] == "DeviceType.CUDA"
|
||
|
|
)
|
||
|
|
|
||
|
|
baseline_timing = summarize(samples["baseline"])
|
||
|
|
candidate_timing = summarize(samples["candidate"])
|
||
|
|
rows.append({
|
||
|
|
"block": index,
|
||
|
|
"candidate_vs_baseline": parity,
|
||
|
|
"baseline_vs_capture": baseline_capture,
|
||
|
|
"candidate_vs_capture": candidate_capture,
|
||
|
|
"baseline": {
|
||
|
|
"timing": baseline_timing,
|
||
|
|
"memory": baseline_memory,
|
||
|
|
"launches": baseline_launches,
|
||
|
|
},
|
||
|
|
"candidate": {
|
||
|
|
"timing": candidate_timing,
|
||
|
|
"memory": candidate_memory,
|
||
|
|
"launches": candidate_launches,
|
||
|
|
},
|
||
|
|
"p50_improvement_percent":
|
||
|
|
(baseline_timing["p50_s"] - candidate_timing["p50_s"])
|
||
|
|
/ baseline_timing["p50_s"]
|
||
|
|
* 100.0,
|
||
|
|
})
|
||
|
|
del adaln_values
|
||
|
|
|
||
|
|
result = {
|
||
|
|
"device": torch.cuda.get_device_name(),
|
||
|
|
"torch": torch.__version__,
|
||
|
|
"workload": workload,
|
||
|
|
"rounds": args.rounds,
|
||
|
|
"blocks": rows,
|
||
|
|
}
|
||
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
args.output.write_text(json.dumps(result, indent=2) + "\n")
|
||
|
|
print(json.dumps(result, indent=2), flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|