File size: 4,246 Bytes
22a73d3 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | #!/usr/bin/env python3
"""Two-prompt 10s768p probe for the direct-route/reweight Spark kernel."""
from __future__ import annotations
import concurrent.futures
import json
from pathlib import Path
import comfyui_sol_4way_benchmark_20260923 as base
import comfyui_spark_sol_10prompt_benchmark_20260924 as source
NAME = "comfyui_spark_fanout16_final_probe_2prompt_10s768p_20260925"
ROOT = Path("/autodl-fs/data/h3_experiments") / NAME
OUT = Path("/autodl-fs/data/h3_outputs") / NAME
BASELINE = Path(
"/autodl-fs/data/h3_experiments/"
"comfyui_sol_tau1_vs_spark_10prompt_10s768p_20260924"
)
GPUS = (0, 1)
PORTS = {0: 8370, 1: 8371}
def graph(case, target, steps):
value = source.denoise_graph("spark_topk10", case, target, steps)
value["2"]["inputs"].update(
ablation_mode="full",
video_tail_mode="dense",
global_anchor_dtype="float32",
)
value["4"]["inputs"].update(width=1344, height=768, length=240)
return value
def run_gpu(gpu, case):
port = PORTS[gpu]
warm = ROOT / "warmup" / f"gpu{gpu}.safetensors"
warm_result = base.queue_and_wait(port, graph(case, warm, 5), source.SAMPLER_NODE)
target = ROOT / "latents" / f"case_{case['index']:02}_{case['sample_id']}.safetensors"
result = base.queue_and_wait(port, graph(case, target, 20), source.SAMPLER_NODE)
return {
"case": case["index"],
"sample_id": case["sample_id"],
"gpu": gpu,
"sampler_seconds": result["sampler_seconds"],
"warmup_seconds": warm_result["sampler_seconds"],
"latent_path": str(target),
"latent_sha256": base.sha256(target),
}
def main():
if ROOT.exists() or OUT.exists():
raise FileExistsError(f"refusing to overwrite {ROOT} or {OUT}")
for path in (ROOT / "warmup", ROOT / "latents", OUT):
path.mkdir(parents=True)
for gpu in GPUS:
for section in ("user", "temp", "input"):
(ROOT / section / f"gpu{gpu}").mkdir(parents=True, exist_ok=True)
cases = source.cases()[:2]
baseline = json.loads((BASELINE / "denoise_summary.json").read_text())
sol_by_case = {
row["case"]: row["sampler_seconds"]
for row in baseline["methods"]["sol_tau1_extra0"]["records"]
}
base.write_json(ROOT / "protocol.json", {
"name": NAME,
"pipeline": "ComfyUI-native MiniMax-H3 denoise-only",
"settings": {"duration": "10s", "resolution": "1344x768", "steps": 20,
"seed": 42, "warmup": "excluded 5-step denoise", "gpus": list(GPUS)},
"spark": {"topk_ratio": 0.1, "reweight": "global", "reblock_reuse_layers": 1,
"route": "official INT8 route + in-kernel exact Top-K",
"reweight_transport": "direct complement bitmap",
"reblock_transport": "permutation fused into Sol preprocess/output scatter",
"reblock_route": "fanout16 exact fused compact route",
"landmark_count": 32, "group_size": 1},
"baseline": str(BASELINE),
"cases": cases,
})
servers = []
try:
base.ROOT = ROOT
for gpu in GPUS:
process, log = base.start_server(f"gpu{gpu}", gpu, PORTS[gpu], OUT)
servers.append((f"gpu{gpu}", process, log))
for gpu in GPUS:
base.wait_server(PORTS[gpu])
with concurrent.futures.ThreadPoolExecutor(max_workers=2) as pool:
rows = list(pool.map(lambda pair: run_gpu(*pair), zip(GPUS, cases, strict=True)))
for row in rows:
row["sol_baseline_seconds"] = sol_by_case[row["case"]]
row["speedup_vs_sol"] = row["sol_baseline_seconds"] / row["sampler_seconds"]
base.write_json(ROOT / "results.json", {
"status": "complete",
"rows": rows,
"mean_spark_seconds": sum(r["sampler_seconds"] for r in rows) / len(rows),
"mean_sol_seconds": sum(r["sol_baseline_seconds"] for r in rows) / len(rows),
"ratio_of_means": (
sum(r["sol_baseline_seconds"] for r in rows)
/ sum(r["sampler_seconds"] for r in rows)
),
})
finally:
base.stop_servers(servers)
if __name__ == "__main__":
main()
|