[Kernel][Hardware][AMD] Bf16 mfma opt for ROCm skinny GEMMs

[amd-hhashemi]

Bf16 mfma opt for ROCm skinny GEMMs

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[tjtanaa]

@amd-hhashemi hi.
Thank you for optimizing bf16. How much perf gain could we expect on mi300?

[amd-hhashemi]

@amd-hhashemi hi. Thank you for optimizing bf16. How much perf gain could we expect on mi300?

Hey, this optimization shows 25% speedup on llama3 bf16 batch-1 on MI300. The prior solution does expensive bf16->float conversion followed by FMA ops. This optimization avoids that by using MFMAs instead, which is much more efficient.

[SageMoore]

Hi, @amd-hhashemi. Thanks for the contribution! Could you just run a quick serving benchmark to make sure there are no obvious perf regressions? I'm somewhat fuzzy on the exact cases that skinny gemm is enabled but I assume that it will be used in llama 3.1 8B.

Additionally, can you post the benchmark you ran that is giving you 25% speedup?

Serving commands:
vllm serve meta-llama/Llama-3.1-8B-Instruct --port 4444 --disable-log-requests
followed by:
python benchmarks/benchmark_serving.py --model meta-llama/Llama-3.1-8B-Instruct --dataset-name sharegpt --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json --ignore-eos --port 4444

[tjtanaa]

Hi @amd-hhashemi Is there a difference between this PR and https://github.com/ROCm/aiter/blob/main/csrc/kernels/custom_kernels.cu ?

[amd-hhashemi]

You mean this PR?
ROCm#520

There is no difference, I wrote that first, then the upstream version was merged to ROCm. So the ROCm one can be dropped and we'll later merge with upstream.

[amd-hhashemi]

Oh sorry I didn't realize you were point to Aiter.
I didn't know it had been pulled into Aiter.
Although it seems to be the original version, before fp8 or bf16 support was added.

[amd-hhashemi]

Hi SageMoore, I will run serving benchmark.
This is what I ran:
python benchmarks/benchmark_latency.py --model /data
/Meta-Llama-3-8B-Instruct --batch-size 1 --dtype bfloat16

[https://github.com/amd-hhashemi/vllm/blob/main/benchmarks/benchmark_latency.py]

Original reported latency: ~1.07sec
After this optimization: ~0.84sec
(it's actually more like ~22% speedup)
[Note: these numbers were actually on a downsized version of MI300, but since it's a compute bottleneck, it should be same on full MI300. I will verify that too]
The skinny gemms get most heavily used with low batch sizes.

[amd-hhashemi]

[corrected, with warmup runs]

I ran the server benchmark before and after the change. There isn't any change on server throughput test (this is expected, skinny GEMMs only show up in low batch count):

Before:

After this code change:

[tjtanaa]

Oh sorry I didn't realize you were point to Aiter. I didn't know it had been pulled into Aiter. Although it seems to be the original version, before fp8 or bf16 support was added.

Will there be plans to integrate this updated kernel into AITER?