[Kernel] Enhance MoE benchmarking & tuning script

[WoosukKwon]

This PR is to enhance the MoE tuning & benchmarking script which is a bit hacky at the moment. Also, the PR enables using multiple GPUs for benchmarking via Ray.

[WoosukKwon]

@pcmoritz The PR is not ready now. I will ping you once it's ready.

[pcmoritz]

Sounds great, thank you :)

[WoosukKwon]

@pcmoritz This PR is ready now. Sorry for the delay.

[pcmoritz]

One small gotcha I was running into while trying this out is that currently fp8 can't be benchmarked with an FP16 checkpoint, e.g.

python benchmark_moe.py --dtype fp8

errors out since mistralai/Mixtral-8x7B-Instruct-v0.1 is FP16. I think what we should do here is

diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py
index 6796ea401..3f3005e20 100644
--- a/benchmarks/kernels/benchmark_moe.py
+++ b/benchmarks/kernels/benchmark_moe.py
@@ -46,6 +46,8 @@ def benchmark_config(
         w2_scale = torch.randn(num_experts, dtype=torch.float32)
         a1_scale = torch.randn(1, dtype=torch.float32)
         a2_scale = torch.randn(1, dtype=torch.float32)
+        w1 = w1.to(torch.float8_e4m3fn)
+        w2 = w2.to(torch.float8_e4m3fn)
 
     input_gating = torch.empty(num_tokens, num_experts, dtype=torch.float32)
 

since FP8 checkpoints are not widely available yet and also for vLLM FP8 we support running FP16 checkpoints in FP8 :)

[WoosukKwon]

@pcmoritz I addressed your comments. PTAL.

[pcmoritz]

Thanks! I've been using the new script to do some tuning for FP8 and it works like a charm, thanks a lot for improving it -- I'll open a PR with the new configs shortly after I have tested the configs!

Btw, in order to get progress bars, I've been using this modification:

from ray.experimental.tqdm_ray import tqdm

and then where we iterate over the configs:

for config in tqdm(search_space):

This will make sure to print progress bars without messing up stdout and it works like this: https://docs.ray.io/en/latest/ray-observability/user-guides/configure-logging.html#distributed-progress-bars-tqdm

Feel free to add it (don't worry it is currently in the experimental namespace -- I think it is one of the APIs that should be stabilized and I'll look into that).

[WoosukKwon]

@pcmoritz Ray tqdm is really cool! I actually wanted to have exactly the same feature. Happy to add that!