[Kernel] Added flashinfer fp8 per-tensor gemms

[nvjullin]

Purpose

Added fp8 gemms from flashinfer.
The added gemms have better or same perf as the original gemms so we use it as the default.
For gemm sizes with small M, the added gemms are marginally faster.
For gemm sizes with large M, the added gemms are much faster.
These are the results for llama3 ISL=OSL=1024 concurrency=128 max_num_batched_tokens=8192 TP1.
As expected, TPOT is roughly the same but TTFT improved by ~13%.

Requires flashinfer autotuning, so depends on #22346 (merged).
Funcionality depends on flashinfer PR flashinfer-ai/flashinfer#1479 (merged).
Perf numbers depends on flashinfer PR flashinfer-ai/flashinfer#1491 (merged).
Requires next flashinfer release including aforementeiond PRs and vllm updating flashinfer version. (updated flashinfer)

old

============ Serving Benchmark Result ============
Successful requests:                     384       
Maximum request concurrency:             128       
Benchmark duration (s):                  133.38    
Total input tokens:                      392638    
Total generated tokens:                  393216    
Request throughput (req/s):              2.88      
Output token throughput (tok/s):         2948.16   
Total Token throughput (tok/s):          5892.00   
---------------Time to First Token----------------
Mean TTFT (ms):                          2955.22   
Median TTFT (ms):                        1330.74   
P99 TTFT (ms):                           10734.06  
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          40.50     
Median TPOT (ms):                        42.68     
P99 TPOT (ms):                           45.28     
---------------Inter-token Latency----------------
Mean ITL (ms):                           40.50     
Median ITL (ms):                         32.10     
P99 ITL (ms):                            653.22    
==================================================

new

============ Serving Benchmark Result ============
Successful requests:                     384       
Maximum request concurrency:             128       
Benchmark duration (s):                  128.33    
Total input tokens:                      392638    
Total generated tokens:                  393216    
Request throughput (req/s):              2.99      
Output token throughput (tok/s):         3064.22   
Total Token throughput (tok/s):          6123.93   
---------------Time to First Token----------------
Mean TTFT (ms):                          2611.05   
Median TTFT (ms):                        1182.85   
P99 TTFT (ms):                           9319.07   
-----Time per Output Token (excl. 1st token)------
Mean TPOT (ms):                          39.19     
Median TPOT (ms):                        41.07     
P99 TPOT (ms):                           44.23     
---------------Inter-token Latency----------------
Mean ITL (ms):                           39.19     
Median ITL (ms):                         31.65     
P99 ITL (ms):                            560.29    
==================================================

lm_eval shows

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value|   |Stderr|
|-----|------:|----------------|-----:|-----------|---|----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.944|±  |0.0103|
|     |       |strict-match    |     5|exact_match|↑  |0.886|±  |0.0142|

Test Plan

Tests to be added.

Test Result

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[gemini-code-assist]

Code Review

This pull request introduces support for FlashInfer's FP8 GEMM kernels, which is expected to improve performance, particularly for large batch sizes. The changes primarily involve refactoring the GEMM dispatch logic to accommodate a new 'flashinfer' backend and adding an optimization to pre-calculate combined scales. While the implementation is largely sound, I've identified a critical issue in the new FlashInfer wrapper where the output tensor is not reshaped, potentially causing shape mismatches for inputs with more than two dimensions.

[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.

🚀

[mgoin]

@nvjullin The Blackwell Test failures look clearly related https://buildkite.com/vllm/ci/builds/27727/steps/canvas?jid=0198c767-0112-49f4-9f26-c9fef601374c#0198c767-0112-49f4-9f26-c9fef601374c/98-3479

[nvjullin]

@mgoin the remaining errors are all something about huggingface gateway timeout

[ProExpertProg]

A few minor notes.

This might be a bit too urgent but in general we should really improve the fp8 scaled_mm dispatching. I started a draft pr #19434 but never got around to it.

[nvpohanh]

@nvjullin Please address the comments and rebase this PR. Thanks!

[nvjullin]

This might be a bit too urgent but in general we should really improve the fp8 scaled_mm dispatching. I started a draft pr #19434 but never got around to it.

@mgoin Is there a specific maintainability problem with the current dispatching code you wish to be addressed?

[ProExpertProg]

Is there a specific maintainability problem with the current dispatching code you wish to be addressed?

No I think it's out of the scope of this PR. But if you look at dispatching for int8 or Marlin/Machete, that's closer to something we want to do in general when we want to dispatch between multiple possible implementations.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @nvjullin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[mgoin]

LGTM to get in, thanks. We should follow up with using an Enum instead of raw strings