[Kernel] change benchmark script so that result can be directly used; tune moe kernel in A100/H100 with tp=2,4,8

[youkaichao]

cc @WoosukKwon

[pcmoritz]

Thanks for the refactorings! While you are touching this code, one thing that would be wonderful to do is to keep track of the timings for the best configuration for each batch size. This could e.g. be done by writing them to a separate file. This would allow you to decide if a new configuration is better than the old one.

Also note that running the script as-is will likely not produce optimal results in some settings, since there is a bunch of parameter pruning going on at the moment (e.g. for the batch size). Sometimes it is important to look at the values found and then expand the search space if it runs into the boundaries :)

[youkaichao]

@pcmoritz This manual kernel tuning is kind of temporary. Going forward, we plan to use triton.autotune to automatically tune these configs. So we don't need to invest too much time here.

[WoosukKwon]

The configuration here is quite different from the one we have right now. Could you compare the old and new ones by benchmarking the end-to-end performance (e.g., using benchmark_throughput.py on Mixtral)?

[WoosukKwon]

Also, it'd be nice if you can benchmark other configs as well, if not all.

[youkaichao]

iirc, LLaMA models don't use MoE. Do you mean mixtral models?

[WoosukKwon]

Yep I mean mixtral, not llama

[youkaichao]

do you have some common setting for me to run the throughput test? Otherwise I'm blindly running python benchmarks/benchmark_throughput.py --input-len 100 --output-len 100. Not sure if input 100 tokens and output 100 tokens are the cases people care the most.

[youkaichao]

I will use python benchmarks/benchmark_throughput.py --model=mistralai/Mixtral-8x7B-Instruct-v0.1 --input-len 1000 --output-len 50 from #2293 (comment) .

[pcmoritz]

@pcmoritz This manual kernel tuning is kind of temporary. Going forward, we plan to use triton.autotune to automatically tune these configs. So we don't need to invest too much time here.

In my experience triton.autotune is far too slow to be useful (unless the configs have already been run / are cached) :)

[youkaichao]

command: python benchmarks/benchmark_throughput.py --model=mistralai/Mixtral-8x7B-Instruct-v0.1 --input-len 1000 --output-len 50 -tp 2

H100 GPU:

benchmark	no config	w/o this PR	w/ PR
tp=2	11.73 requests/s, 12313.00 tokens/s	14.11 requests/s, 14816.20 tokens/s	15.69 requests/s, 16477.71 tokens/s
tp=4	18.07 requests/s, 18974.30 tokens/s	same as no config	22.06 requests/s, 23166.51 tokens/s
tp=8	24.80 requests/s, 26039.89 tokens/s	same as no config	28.17 requests/s, 29577.75 tokens/s

A100 GPU: TODO (don't have 8*A100 GPU at hand now)

@WoosukKwon benchmarking results are quite promising!

[youkaichao]

In my experience triton.autotune is far too slow to be useful (unless the configs have already been run / are cached) :)

Will definitely try to cache tuned configs!

[WoosukKwon]

@youkaichao Awesome! Could you 1) update the PR with the current main and 2) fix the lint error by running ./format.sh? You will have to run pip install -r requirements-dev.txt before it.

[youkaichao]

@WoosukKwon lint is good now 👌

[WoosukKwon]

LGTM! Thanks for the PR! Excited about the performance improvement!