From 051e40e82b71c81c99a86190da039830ba0713b5 Mon Sep 17 00:00:00 2001 From: sazc Date: Mon, 24 Feb 2025 10:59:22 +0800 Subject: [PATCH] tests: Triton had remove the fast_flush parameter from do_bench (#4485) --- tests/test_flash_mla.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_flash_mla.py b/tests/test_flash_mla.py index a057f6a..9c5cd90 100644 --- a/tests/test_flash_mla.py +++ b/tests/test_flash_mla.py @@ -87,7 +87,7 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, causal, varlen): cal_diff(out_flash, out_torch, "out") cal_diff(lse_flash, lse_torch, "lse") - t = triton.testing.do_bench(flash_mla, fast_flush=False) + t = triton.testing.do_bench(flash_mla) FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2 bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d + b * s_q * h_q * dv) * (torch.finfo(dtype).bits // 8) print(f"{t:.3f} ms, {FLOPS / 10 ** 9 / t:.0f} TFLOPS, {bytes / 10 ** 6 / t:.0f} GB/s")