Edit bechmark_with_mem.py

Attention-Benchmarking/bechmark_with_mem.py

@@ -114,7 +114,7 @@ for causal in causal_vals:
 
             results[config]["Pytorch"] = {"fwd_time": f, "bwd_time": b, "mem_used_MB": mem_used, "mem_peak_MB": mem_peak}
 
-            # xFormers
+            # FlashAttention mit xFormers
             if xops is not None:
                 q, k, v = [torch.randn(batch_size, seqlen, nheads, headdim, device=device, dtype=dtype, requires_grad=True) for _ in range(3)]
                 f, b, mem_used, mem_peak = benchmark_with_memory(