Put a real enum.

enum KvDtype{
    Fp8(Path)
}

...
#[clap(long, env, value_enum)]
Option<KvDtype>

This should work much better. None is equivalent to auto (it just means the user hasn't specified anything we can do whateverwe want with it).

KvDtype will automatically be sanitized/error checked (String isn't since all strings are available).

I tried putting Fp8(Path) directly in clap, I'm not sure it actually works in clap internals but this is what we want, if fp8 is chosen we need a path for the scales. and Path should also ensure the string is a valid path.

Maybe clap doesn't support algebraic enunms and we can't have Fp8(Path) and need Fp8 instead.
In that case you need to handle validation early (There are other forms of validation in that layer, before pushing the args to the shard).

All CLI validation should happen here, as early as possible, with the best possible error messages.

Why an external file is needed ? Can't the user specify a value ?
If this is linked per model, shouldn't the model config/repo contain that information ?

If it's needed with kvcache=fp8 let's try to make sure it's actually. Ideally it's one option for users, if not possible we need manual validation here (and vaildation can be skipped later)

(I'm trying to avoid adding too many flags, TGI already has too many, and since we don't break, we never remove stuff that was added, that's why if we can read the information from some consistant config in the repo it keeps the interface for the user simpler)

Updated to load the scales from the checkpoint!

This should probably be done higher in the stack (ideally in launcher directly).

Rust is much more efficient at running these kind of checks but most importantly errors should happen as early as possible (and launcher has all the user flags too).

Why are you not fusing this with the first validator ?

I'm wondering what kind of bug this is supposed to prevent.

If the scaling are contained directly within a single repo (so not user supplied) the validity should be obvious (no need for extra keys).

If it is user supplied, well it is unlikely to contain a model_type, no ?

Optional[str] without validation is acceptable (since the launcher is responsible for validation already).

Can we avoid passing 2 objects all the time (they seem highly interlinked, so we could probably fuse them in a single object,no ?)

Done

Huge no-no.

We really don't like mutating self state of models (it makes reasoning about who modified your values (or not) much harder.

You need to be sending the scales directly at load time (possibly even transparent through config or weights)

We do the same for quantize and speculator for instance.

Codewise it touches a lot of places, but these are about to be factored out.

Done, the params should be created at load time

Remove all these, state should already be correct.

Done

And why does that imply that rocm needs a factor of 2 ?

If the reason is that much global, shouldn't it be handled directly on load?

@Narsil I believe it is related to https://onnx.ai/onnx/technical/float8.html and the diff between e4m3fn and e4m3 (different exponent bias). Is that so @mht-sharma?

But shouldn't it be based on this param https://github.com/vllm-project/vllm/blob/319ad7f1d386699e94f629341c9988a926821f24/tests/fp8_kv/llama2-7b-fp8-kv/kv_cache_scales.json#L4 ? Also, we have to kind of decide whether we want to follow the scheme in vllm to store quantization params in a json and keep compatibility for e.g. models on the Hub, or not.

Comment from AMD in vllm:

We do *2 only for HIP, to deal with the difference in numeric from our chip. after *2 overall effect is identical as without it on NV.

I will add this a comment in code

Yes, this is the difference between e4m3fn and e4m3 formats. If using a json like the one liked, there should be a check on the dtype.

Not needed if done at load time.

Fused?

If done correctly, you don't need any of that because you handle this at load time (and every model will need to be updated to support the new values)

This should be handled all the way back in the launcher (again ideally directly by clap, if not possible manually in the rust part).

And it should be hard error, not a soft one (paramters sent by the user don't make sense, we never silently ignore.)

Fused?

Everything here should be hard error (I think the standard ones would do fine).

If a user sends invalid information, we shouldn't silently ignore. They should fix it.

This one isn't necessary here, it should already be inferrable from kv_cache[0].

to be precise, on Instinct I think it is E4M3FN that is used

I would use builtin dataclasses and post_init here but it is a matter of taste

It would be worth to explain what is FP8 KV Cache. Readers may not be familiar with it (does it mean attention computation is in fp8? etc)

enhancing inference speed on both Nvidia and AMD GPUs. This feature significantly boosts performance and memory efficiency, enabling faster and more scalable text generation.

This is kind of vague.

we can greatly reduce the memory footprint. This reduction allows for improved throughput in text generation tasks

It would be worth IMO to show numbers / a chart here to get a grasp of what greatly means, in which case there is indeed a speedup, etc

We have a couple of options for numbers and charts

• Max Batch Total Tokens: The logs display the max batch total tokens, which increase when using the FP8 KV cache. We could create a chart showing the max batch total tokens in both cases (with and without the FP8 KV cache).

• Throughput: Currently, I have created a custom script using AsyncClient to send 500 requests simultaneously with asyncio.gather. This provides a rough estimate of throughput. @Narsil , do you have any suggestions on calculating throughput more precisely?

Technically, Instinct supports E4M3FNUZ, not E4M3. https://onnx.ai/onnx/technical/float8.html

Let's maybe put a full runnable command with docker run etc? Good inspiration could be https://huggingface.co/docs/transformers/main/en/perf_infer_gpu_one?install=NVIDIA#flashattention-2

same + maybe have a model_id which indeed has scaling factors here.

How do we know whether they are for E4M3FNUZ or E4M3FN format?

I guess the current tools Nvidia AMMO and AutoFP8 uses E4M3FN. Currently there is no flag to determine the format (unless ofcourse the weight is quantized). But I could add a check and add a parameter for this in the checkpoint.

Updated reading this from config. See: a7909e6

Todo

Needs a description

Things are a bit harder to read if dtype attribute is used to mean the KV cache storage pytorch dtype. For some other models (gemma, idefics), self.dtype is used with an other meaning, being the weights dtype.

How do we know whether serialized scales are in E4M3FN or E4M3FNUZ format? I think depending on that, the logic should be different here.

Updated the code to read the corresponding dtype as kv_cache_torch_dtype from config.

Added the format in the README.md and utility script to add kv_cache_torch_dtype when quantising model

See: a7909e6

If we target vllm/tgi intercompatibility, why couldn't we load directly e.g. neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV in TGI? Is it not loadable in vllm?

Also, by quantized FP8 model, do you mean a model whose weights are quantized to fp8? How does it relate to FP8 KV cache? To me to obtain the KV cache scales you would simply need to have calibration data passing through the network & collecting stats on the KV cache.

It feels like the the KV cache scales from a quantized model (like neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV), whose KV cache scales may have been obtained after the weights quantization (?), may have an unnecessary bias due to being inferred from calibration on the quantized model, not the unquantized one.

The model neuralmagic/Meta-Llama-3-8B-Instruct-FP8-KV has both the weights quantized along with KV scales. However, I don't think we can load the FP8 weights in TGI yet. And also one would run an FP16 model with FP8 KV cache so we need such a checkpoint

You are right; it may have an additional bias due to calibration after weight quantization. The bias might have been low, so I couldn't find a noticeable difference during inference.

I have tested two other models generated using scales from Nvidia AMMO (VLLM uses this also), which may not have this bias. The quantizer can accept the quantized format as full precision and provide the model with KV scales. We can extract scales from that.

Here is the link to the models: Llama-2-70b-chat-hf-FP8-KV

I will add these models and provide an example for this.

Updated the example with Nvidia AMMO

Do you have an example of such model?

Updated and added a model to readme.

I think it would be worth to have some evaluation metrics for an example model

is this a typo of "baichuan"?

Yes. Thanks for pointing it out

Is it should be "ROCm uses FP8 format with fp8_e4m3fnuz, whereas NVIDIA GPU uses fp8_e4m3fn"?

Yes that's correct