Support calibrating kv cache scales

[mgoin]

Adds a kv_cache_quant_targets quant config argument that attaches output_scales to the specified Linear modules. This means we will end up with k_proj.output_scale and v_proj.output_scale after activation calibration. For the final checkpoint, we add a pass to take the maximum of k_proj.output_scale and v_proj.output_scale, and place the result in the parent of those modules (the Attention module) as a single kv_scale, which is needed to match the representation in vLLM.

Also includes a decent chunk of refactoring to allow for no examples to be passed in for weight quantization, renaming for clearer understanding of modules, making "re:.*lm_head" not a required ignored pattern but just a default, and disabling torch._scaled_mm for easier usage on CPU.

A new example is included to show how to enable this functionality

from datasets import load_dataset
from transformers import AutoTokenizer

from auto_fp8 import AutoFP8ForCausalLM, BaseQuantizeConfig

pretrained_model_dir = "meta-llama/Meta-Llama-3-8B-Instruct"
quantized_model_dir = "Meta-Llama-3-8B-Instruct-FP8-KV"

tokenizer = AutoTokenizer.from_pretrained(pretrained_model_dir, use_fast=True)
tokenizer.pad_token = tokenizer.eos_token

ds = load_dataset("mgoin/ultrachat_2k", split="train_sft").select(range(512))
examples = [tokenizer.apply_chat_template(batch["messages"], tokenize=False) for batch in ds]
examples = tokenizer(examples, padding=True, truncation=True, return_tensors="pt").to("cuda")

quantize_config = BaseQuantizeConfig(
    quant_method="fp8",
    activation_scheme="static",
    ignore_patterns=["re:.*lm_head"],
    kv_cache_quant_targets=("k_proj", "v_proj"),
)

model = AutoFP8ForCausalLM.from_pretrained(pretrained_model_dir, quantize_config)
model.quantize(examples)
model.save_quantized(quantized_model_dir)