Add files using upload-large-folder tool

gpu_run.py

@@ -0,0 +1,126 @@
+# gpu_run.py
+# 用法：
+#   python gpu_run.py
+#   python gpu_run.py --target 0.6 --duration 0
+#   CUDA_VISIBLE_DEVICES=0,1 python gpu_run.py --target 0.6 --size 4096 --streams 4
+#
+# 说明：
+# - 目标是把每张可见 GPU 的平均利用率维持在大约 target（默认 60%）
+# - 原理：每个周期里“忙一会儿 + 休一会儿”，让 nvidia-smi 统计窗口看到约 60% 的利用率
+# - 不同显卡/驱动/CUDA 版本会有偏差；要更接近 60%，优先调 --size 和 --cycle-ms
+
+import argparse
+import threading
+import time
+import torch
+
+
+def worker(
+    gpu_id: int,
+    target: float,
+    duration: float,
+    size: int,
+    streams: int,
+    cycle_ms: int,
+):
+    device = f"cuda:{gpu_id}"
+    torch.cuda.set_device(device)
+    torch.backends.cuda.matmul.allow_tf32 = True
+    torch.set_float32_matmul_precision("high")
+
+    # 用 FP16/BF16 大矩阵乘法制造稳定算力负载
+    # 如果你的卡不太适合 fp16，也可改成 torch.float32
+    dtype = torch.float16
+    if torch.cuda.get_device_capability(gpu_id)[0] >= 8:
+        # Ampere/Hopper 往往 BF16 也很稳
+        dtype = torch.bfloat16
+
+    # 为每个 stream 预分配张量，避免循环中频繁申请显存
+    cuda_streams = [torch.cuda.Stream(device=device) for _ in range(streams)]
+    tensors = []
+    for _ in range(streams):
+        a = torch.randn((size, size), device=device, dtype=dtype)
+        b = torch.randn((size, size), device=device, dtype=dtype)
+        tensors.append((a, b))
+
+    # 预热
+    for s, (a, b) in zip(cuda_streams, tensors):
+        with torch.cuda.stream(s):
+            for _ in range(4):
+                _ = a @ b
+    torch.cuda.synchronize(device)
+
+    busy_ms = max(1, int(cycle_ms * target))
+    idle_ms = max(0, cycle_ms - busy_ms)
+
+    start = time.time()
+    end_time = start + duration if duration > 0 else float("inf")
+
+    print(
+        f"[GPU {gpu_id}] start | target={target:.2f}, size={size}, streams={streams}, "
+        f"cycle={cycle_ms}ms, busy={busy_ms}ms, idle={idle_ms}ms, dtype={dtype}"
+    )
+
+    while time.time() < end_time:
+        t0 = time.time()
+
+        # busy phase：在设定时间内持续发 matmul
+        while (time.time() - t0) * 1000 < busy_ms:
+            for s, (a, b) in zip(cuda_streams, tensors):
+                with torch.cuda.stream(s):
+                    _ = a @ b
+
+        # 确保本轮计算真正执行完，避免堆积过多 kernel
+        torch.cuda.synchronize(device)
+
+        # idle phase：CPU sleep，让 GPU 闲下来，形成平均 60% 利用率
+        if idle_ms > 0:
+            time.sleep(idle_ms / 1000.0)
+
+    torch.cuda.synchronize(device)
+    print(f"[GPU {gpu_id}] done")
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--target", type=float, default=0.6, help="目标平均利用率，默认 0.6")
+    parser.add_argument("--duration", type=float, default=0, help="运行秒数；0 表示一直跑")
+    parser.add_argument("--size", type=int, default=4096, help="矩阵边长，默认 4096")
+    parser.add_argument("--streams", type=int, default=4, help="每张卡并发 stream 数，默认 4")
+    parser.add_argument("--cycle-ms", type=int, default=1000, help="占空比周期，默认 1000ms")
+    args = parser.parse_args()
+
+    if not torch.cuda.is_available():
+        raise RuntimeError("没有检测到 CUDA GPU。")
+
+    if not (0 < args.target <= 1.0):
+        raise ValueError("--target 必须在 (0, 1] 内")
+
+    num_gpus = torch.cuda.device_count()
+    print(f"Detected {num_gpus} visible GPU(s)")
+    for i in range(num_gpus):
+        print(f"  GPU {i}: {torch.cuda.get_device_name(i)}")
+
+    threads = []
+    for gpu_id in range(num_gpus):
+        t = threading.Thread(
+            target=worker,
+            args=(
+                gpu_id,
+                args.target,
+                args.duration,
+                args.size,
+                args.streams,
+                args.cycle_ms,
+            ),
+            daemon=False,
+        )
+        t.start()
+        threads.append(t)
+
+    for t in threads:
+        t.join()
+
+
+if __name__ == "__main__":
+    main()