From 0e45a22094ad378740def75a34e2d097b9347b34 Mon Sep 17 00:00:00 2001 From: princepride Date: Fri, 29 May 2026 02:59:58 +0000 Subject: [PATCH] fix pre commit Signed-off-by: princepride --- benchmarks/kernels/mot_linear_benchmarks.py | 10 +++++----- tests/diffusion/kernels/mot/test_mot_linear.py | 8 ++++---- tests/diffusion/kernels/mot/test_mot_norm.py | 4 ++-- 3 files changed, 11 insertions(+), 11 deletions(-) diff --git a/benchmarks/kernels/mot_linear_benchmarks.py b/benchmarks/kernels/mot_linear_benchmarks.py index 7d22c8209..20674272e 100644 --- a/benchmarks/kernels/mot_linear_benchmarks.py +++ b/benchmarks/kernels/mot_linear_benchmarks.py @@ -64,7 +64,7 @@ def clear_triton_cache(): """Clear Triton JIT compilation cache and Python/CUDA memory.""" gc.collect() if torch.cuda.is_available(): - torch.cuda.empty_cache() + torch.accelerator.empty_cache() try: if hasattr(triton, "runtime") and hasattr(triton.runtime, "cache") and hasattr(triton.runtime.cache, "clear"): triton.runtime.cache.clear() @@ -487,18 +487,18 @@ def benchmark_config( # JIT warmup run() - torch.cuda.synchronize() + torch.accelerator.synchronize() # Capture 1 invocations with CUDA Graph graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): run() - torch.cuda.synchronize() + torch.accelerator.synchronize() # Warmup replays for _ in range(5): graph.replay() - torch.cuda.synchronize() + torch.accelerator.synchronize() start_event = torch.cuda.Event(enable_timing=True) end_event = torch.cuda.Event(enable_timing=True) @@ -507,7 +507,7 @@ def benchmark_config( for _ in range(num_iters): if cache_flusher is not None: cache_flusher.zero_() - torch.cuda.synchronize() + torch.accelerator.synchronize() start_event.record() graph.replay() diff --git a/tests/diffusion/kernels/mot/test_mot_linear.py b/tests/diffusion/kernels/mot/test_mot_linear.py index 4f221e9e2..ef9ca4237 100644 --- a/tests/diffusion/kernels/mot/test_mot_linear.py +++ b/tests/diffusion/kernels/mot/test_mot_linear.py @@ -176,21 +176,21 @@ def _benchmark(fn, warmup: int = 20, iters: int = 100) -> float: cache_flusher = torch.empty(int(256 * 1024 * 1024 / 4), dtype=torch.int32, device="cuda") for _ in range(warmup): fn() - torch.cuda.synchronize() + torch.accelerator.synchronize() t_flush_start = time.perf_counter() for _ in range(iters): cache_flusher.zero_() - torch.cuda.synchronize() + torch.accelerator.synchronize() flush_time_total = time.perf_counter() - t_flush_start # 3. Measure the total time of "flush cache + operator execution" - torch.cuda.synchronize() + torch.accelerator.synchronize() t_total_start = time.perf_counter() for _ in range(iters): cache_flusher.zero_() fn() - torch.cuda.synchronize() + torch.accelerator.synchronize() total_time = time.perf_counter() - t_total_start # 4. Asynchronous subtraction separation diff --git a/tests/diffusion/kernels/mot/test_mot_norm.py b/tests/diffusion/kernels/mot/test_mot_norm.py index 736204eec..713e13885 100644 --- a/tests/diffusion/kernels/mot/test_mot_norm.py +++ b/tests/diffusion/kernels/mot/test_mot_norm.py @@ -51,11 +51,11 @@ def _benchmark(fn, warmup: int = 20, iters: int = 200) -> float: """Return mean latency in milliseconds.""" for _ in range(warmup): fn() - torch.cuda.synchronize() + torch.accelerator.synchronize() t0 = time.perf_counter() for _ in range(iters): fn() - torch.cuda.synchronize() + torch.accelerator.synchronize() return (time.perf_counter() - t0) / iters * 1000.0