[Usage] How to do KV cache quantization?

[CharlesRiggins]

I followed this example Activation quantization to fp8 and got an FP8 quantized model. I also want to run the model with FP8 E4M3 KV cache. So my question is how do I set the kv_cache_scheme when I quantize the model?

[mgoin]

@CharlesRiggins I am working on this right now and will share the PR for the example later today.

For a quick example, it is just a new entry inside of QuantizationModifier so you can add it to a recipe like this:

recipe = """
quant_stage:
    quant_modifiers:
        QuantizationModifier:
            ignore: ["lm_head"]
            config_groups:
                group_0:
                    weights:
                        num_bits: 8
                        type: float
                        strategy: tensor
                        dynamic: false
                        symmetric: true
                    input_activations:
                        num_bits: 8
                        type: float
                        strategy: tensor
                        dynamic: false
                        symmetric: true
                    targets: ["Linear"]
            kv_cache_scheme:
                num_bits: 8
                type: float
                strategy: tensor
                dynamic: false
                symmetric: true
"""

[mgoin]

Take a look here: #113

[CharlesRiggins]

Great. I tried the example and it worked. Thank you!