[Quantization][FP8] Add support for FP8 models with input_scale for output projection and QK quantization

[rasmith]

This PR adds fp8 quantization support for:

1. Preserving FP8 quantization after FA output, so that output of FA into the next layer will be FP8
   • this uses the model parameter self_attn.q/k/v_proj.input_scale and passes it as _out_scale to the FA kernel
2. During execution of FA loop, the quantity softmax(QK^T) can also be quantized as FP8
   • this uses the model parameter self_attn.prob_output_scale

These are using Quark quantized models, in particular using Llama-3.1-8B-Instruct-FP8-QKV-Prob

So far, only Llama is being supported, but more can be added later.

[tlrmchlsmth]

Let’s make sure this doesn’t break V1 before landing

[rasmith]

Does this only work with vLLM V0? Ideally this would get added to V1 (and then it wouldn't necessarily need to happen in V0 since we plan to deprecate it "soon").

What will happen if someone tries to run such a model in V1?

BTW is amd/Llama-3.1-8B-Instruct-FP8-KV the right huggingface stub to try this out?

It's not V0 ready, but I think it was agreed that a follow-on PR would be OK. You can use amd/Llama-3.1-8B-Instruct-FP8-KV or amd/Llama-3.1-70B-Instruct-FP8-QKV-Prob. Or if the model has 'k_proj/v_proj/k_proj.output_scaleand is set forfp8orquark` quantization.

[rasmith]

Let’s make sure this doesn’t break V1 before landing

OK, I added a check in arg_utils to check if fp8 quantization being used and it just calls _raise_or_fallback.

[rasmith]

@tlrmchlsmth Please take another look? Ran into some problem with VllmConfig._get_quantization_config, it modifies the model_config object that gets passed to it.

However, when running with VLLM_USE_V1=1 there is a check in arg_utils.py that will check for compatibility. The check works for both cases of V0 and V1 now. So, if using V0, everything works, and if using VLLM_USE_V1=1 the check happens and user is just told that fp8/quark fp8 is not implemented for ROCm.

[tlrmchlsmth]

@rasmith could you merge in latest main? It should fix the mamba-ssm