Skip to content

Commit e11880d

Browse files
authored
[Bugfix] Remove triton do_bench fast_flush arg (#16256)
Signed-off-by: Kebe <mail@kebe7jun.com>
1 parent 9351f91 commit e11880d

File tree

1 file changed

+1
-1
lines changed

1 file changed

+1
-1
lines changed

tests/kernels/test_flashmla.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -124,7 +124,7 @@ def ref_mla():
124124
cal_diff(out_flash, out_torch, "out")
125125
cal_diff(lse_flash, lse_torch, "lse")
126126

127-
t = triton.testing.do_bench(flash_mla, fast_flush=False)
127+
t = triton.testing.do_bench(flash_mla)
128128
FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2
129129
bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d +
130130
b * s_q * h_q * dv) * (torch.finfo(dtype).bits // 8)

0 commit comments

Comments
 (0)