[Bugfix] fix rope error when load models with different dtypes

[jinzhen-lin]

Currently, if we load models with different dtypes in the same process, we would get an error like

File ~/.miniconda3/lib/python3.8/site-packages/vllm/_custom_ops.py:89, in rotary_embedding(positions, query, key, head_size, cos_sin_cache, is_neox)
     81 def rotary_embedding(
     82     positions: torch.Tensor,
     83     query: torch.Tensor,
   (...)
     87     is_neox: bool,
     88 ) -> None:
---> 89     vllm_ops.rotary_embedding(positions, query, key, head_size, cos_sin_cache,
     90                               is_neox)

RuntimeError: expected scalar type BFloat16 but found Half

To reproduce:

import torch
from vllm import LLM

model_fp16 = LLM("Qwen/Qwen1.5-0.5B", dtype=torch.half, gpu_memory_utilization=0.4)
model_bf16 = LLM("Qwen/Qwen1.5-0.5B", dtype=torch.bfloat16, gpu_memory_utilization=0.4)

The bug is caused by the rope cache, different dtypes share the same rope module. This PR add dtype to cache key to fix this bug.

[Yard1]

can we pass the dtype as an argument instead?

[jinzhen-lin]

done.

[rkooo567]

QQ: is it difficult to always require to pass the dtype instead?

[jinzhen-lin]

I notice that linear module in vllm set param_dtype as an optional argument, so I think it may be better to keep the same.

[rkooo567]

import torch
from vllm import LLM

model_fp16 = LLM("Qwen/Qwen1.5-0.5B", dtype=torch.half, gpu_memory_utilization=0.4)
model_bf16 = LLM("Qwen/Qwen1.5-0.5B", dtype=torch.bfloat16, gpu_memory_utilization=0.4)

Can you add this as a regression test? And then it lgtm

[jinzhen-lin]

import torch
from vllm import LLM

model_fp16 = LLM("Qwen/Qwen1.5-0.5B", dtype=torch.half, gpu_memory_utilization=0.4)
model_bf16 = LLM("Qwen/Qwen1.5-0.5B", dtype=torch.bfloat16, gpu_memory_utilization=0.4)

Can you add this as a regression test? And then it lgtm

I add a rope module cache test instead of model test, is that ok?

[rkooo567]

Yeah test lgtm!