Support FP8 model fallback KVCache to bfloat16

[changwangss]

I plan to load fp8 model with the following config, Linear is fp8 and kvcache and others op are bf16.

FP8Config(allowlist={"types": ["Linear"], "names": []}, blocklist=blocklist =  {"types": [], "names": []})

when use run_generation.py do model.generate, the error raised.

    outputs = self.model(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1556, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1606, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/changwang/workspace/vllm/optimum-habana-fork/optimum/habana/transformers/models/llama/modeling_llama.py", line 1278, in forward
    layer_outputs = decoder_layer(
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1556, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/torch/nn/modules/module.py", line 1606, in _call_impl
    result = forward_call(*args, **kwargs)
  File "/home/changwang/workspace/vllm/optimum-habana-fork/optimum/habana/transformers/models/llama/modeling_llama.py", line 962, in forward
    hidden_states, self_attn_weights, present_key_value = self.pre_attn(
  File "/home/changwang/workspace/vllm/optimum-habana-fork/optimum/habana/transformers/models/llama/modeling_llama.py", line 1019, in pre_attn
    hidden_states, attn_weights, present_key_value = self.self_attn.pre_attn_forward(
  File "/home/changwang/workspace/vllm/optimum-habana-fork/optimum/habana/transformers/models/llama/modeling_llama.py", line 682, in pre_attn_forward
    key_states = self.k_cache.update(past_key_value[0], key_states, 2, token_idx, self.inp_seq_len)
  File "/home/changwang/workspace/vllm/optimum-habana-fork/optimum/habana/transformers/models/llama/modeling_llama.py", line 426, in update
    prev.index_copy_(dim, idx - 1, cur)
RuntimeError: index_copy_(): self and source expected to have the same dtype, but got (self) Float8_e4m3fn and (source) BFloat16

[HolyFalafel]

Why not use the function?
The default value is:
self.k_proj.weight.dtype

[changwangss]

for recipes FP8Config(allowlist={"types": ["Linear"], "names": []}, blocklist=blocklist = {"types": [], "names": []}), self.k_proj.weight.dtype is torch.float8_e4m3fn, but the past_key dtype should be torch.bfloat16

[HolyFalafel]

same