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

[schoi-habana]

in the init(), self.k_proj is initialized as a KVCache(). isn't isinstance(self.k_cache, KVCache) always true? when does it hit the else part?

[changwangss]

the self.k_proj is the kind of linear,
no, self.k_cache is KVCache for bf16, self.k_cache is PatchedKVCache for fp8.
Scenario: self.k_proj's weight is fp8, but the past_key_value need use bf16 datatype. ( for example,
When the accuracy does not reach 1% loss, we will selectively fallback operator, such as kv_cache.
When I want to load neuralmagic model(https://huggingface.co/neuralmagic/Meta-Llama-3.1-8B-Instruct-FP8), there are only linear fp8 weights in the checkpoints, but no kv cache scale.)

[astachowiczhabana]

@changwangss can you please resolve conflicts?