[Core][Kernels] Enable FP8 KV Cache with Flashinfer backend. + BugFix for kv_cache_dtype=auto

[pavanimajety]

Previous PR with dicussions: #7798

Revert PR due to bug with --kv-cache-dtype=auto : #7982

--kv-cache-dtype=fp8

model_str="neuralmagic/Meta-Llama-3-8B-Instruct-FP8"
model = LLM(model=model_str, quantization="fp8",kv_cache_dtype="fp8")
prompts = [
    "Hello, my name is",
    "The president of the United States is",
    "The capital of France is",
    "New york times is politically sided to ",
]

Results:

$ VLLM_ATTENTION_BACKEND=FLASHINFER /bin/python3 /workspace/vllm_github/test_llm.py
..........

 Prompt: 'Hello, my name is', 
Generated text: ' [Your Name]. I am a [Your Profession/Student] and I am', 
token_ids: array('l', [510, 7927, 4076, 948, 358, 1097, 264, 510, 7927, 50311, 14, 14428, 60, 323, 358, 1097])


 Prompt: 'The president of the United States is', 
Generated text: ' the head of state and the head of government of the United States. The president', 
token_ids: array('l', [279, 2010, 315, 1614, 323, 279, 2010, 315, 3109, 315, 279, 3723, 4273, 13, 578, 4872])


 Prompt: 'The capital of France is', 
Generated text: ' Paris, which is located in the north-central part of the country. Paris is', 
token_ids: array('l', [12366, 11, 902, 374, 7559, 304, 279, 10411, 85181, 961, 315, 279, 3224, 13, 12366, 374])


 Prompt: 'New york times is politically sided to ', 
Generated text: '1 side\nThe New York Times is often accused of being biased towards the left', 
token_ids: array('l', [16, 3185, 198, 791, 1561, 4356, 8691, 374, 3629, 13487, 315, 1694, 48761, 7119, 279, 2163])

--kv-cache-dtype=auto

model_str="neuralmagic/Meta-Llama-3-8B-Instruct-FP8"
model = LLM(model=model_str, quantization="fp8",kv_cache_dtype="auto")

Results:

/bin/python3 /workspace/vllm_github/test_llm.py

........

Processed prompts: 100%|███████████████████████████████████████████████████████████████████| 5/5 [00:00<00:00, 36.28it/s, est. speed input: 254.03 toks/s, output: 580.62 toks/s]


 Prompt: 'Hello, my name is', 
Generated text: ' [Name]. I am a [Your Profession/Student] and I am here', 
token_ids: array('l', [510, 678, 948, 358, 1097, 264, 510, 7927, 50311, 14, 14428, 60, 323, 358, 1097, 1618])


 Prompt: 'The president of the United States is', 
Generated text: ' the head of state and head of government of the United States. The president is', 
token_ids: array('l', [279, 2010, 315, 1614, 323, 2010, 315, 3109, 315, 279, 3723, 4273, 13, 578, 4872, 374])


 Prompt: 'The capital of France is', 
Generated text: ' Paris, which is located in the north-central part of the country. Paris is', 
token_ids: array('l', [12366, 11, 902, 374, 7559, 304, 279, 10411, 85181, 961, 315, 279, 3224, 13, 12366, 374])


 Prompt: 'New york times is politically sided to ', 
Generated text: '1 side\nThe New York Times is often accused of being biased towards the left', 
token_ids: array('l', [16, 3185, 198, 791, 1561, 4356, 8691, 374, 3629, 13487, 315, 1694, 48761, 7119, 279, 2163])

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.
Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which consists a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of default ones by unblocking the steps in your fast-check build on Buildkite UI.

Once the PR is approved and ready to go, please make sure to run full CI as it is required to merge (or just use auto-merge).

To run full CI, you can do one of these:

• Comment /ready on the PR
• Add ready label to the PR
• Enable auto-merge.

🚀

[pavanimajety]

@simon-mo @youkaichao @comaniac Can one of you please take a look?

Thanks!

[pavanimajety]

Can you please help with the CI here? It seems to be yellow without logs

[pavanimajety]

PR is ready for merge :-)

[mgoin]

LGTM thanks for the kernel tests! Just one question about the data_type arg to flashinfer

[mgoin]

Why is this not needed anymore?

[pavanimajety]

The wrappers are throwing dispatch errors with Float8 datatype with the current flashinfer versions(0.1.4 and 0.1.5). I have a couple of future MRs planned. I will enable it if future releases of Flashinfer fixes the dispatch api. Seems like there would be api name changes with flashinfer too.

[mgoin]

I see, this is sufficient for now and it seems we will likely need to address this in flashinfer first

[pavanimajety]

Follow-up: #8009