A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 14:45 UTC
[Bugfix] Remove triton do_bench fast_flush arg (#16256)
Signed-off-by: Kebe <mail@kebe7jun.com>
Kebe committed
Apr 8, 2025 at 21:51 UTC
e11880deea8a181537a84d076730e0f0dd493f0a
1 file changed
+1
-1
tests/kernels/test_flashmla.py
+1
-1
@@ -124,7 +124,7 @@ def test_flash_mla(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size, causal,
124
cal_diff(out_flash, out_torch, "out")
125
cal_diff(lse_flash, lse_torch, "lse")
126
127
- t = triton.testing.do_bench(flash_mla, fast_flush=False)
127
+ t = triton.testing.do_bench(flash_mla)
128
FLOPS = s_q * total_seqlens * h_q * (d + dv) * 2
129
bytes = (total_seqlens * h_kv * d + b * s_q * h_q * d +
130
b * s_q * h_q * dv) * (torch.finfo(dtype).bits // 8)