Can this model run on the A100/sm80 architecture?

#3
by djakip - opened

Anyone tried this yet?

I successfully ran the program on 8*a100 40g. However, after installing VLLM using the code provided by the author, it is necessary to modify the “vllm/_custom_ops.py” file to enable support for kv_cache_dtype.

I successfully ran the program on 8*a100 40g. However, after installing VLLM using the code provided by the author, it is necessary to modify the “vllm/_custom_ops.py” file to enable support for kv_cache_dtype.

Thanks for your reply. Could you share your patch code? I would really appreciate it.

I successfully ran the program on 8*a100 40g. However, after installing VLLM using the code provided by the author, it is necessary to modify the “vllm/_custom_ops.py” file to enable support for kv_cache_dtype.

Hi, thanks for sharing! I'm trying to reproduce your setup. Could you share the specific changes you made to for support?

Find /usr/local/lib/python3.12/dist-packages/vllm/_custom_ops.py

And edit the file

import torch

# ... Find the reshape_and_cache_flash function in the file and replace it with the following content...
def reshape_and_cache_flash(
    key: torch.Tensor,
    value: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    kv_cache_dtype: str,
    k_scale: torch.Tensor,
    v_scale: torch.Tensor,
) -> None:
    # ==========================================
    # Patch for A100 (sm80) FP8 KV cache support
    # ==========================================
    if kv_cache_dtype in ("fp8", "fp8_e4m3"):
        # Convert the input to a safe allocated type (such as bf16) and set to "auto" to bypass the check
        key = key.to(key_cache.dtype)
        value = value.to(value_cache.dtype)
        kv_cache_dtype = "auto"

    torch.ops._C_cache_ops.reshape_and_cache_flash(
        key,
        value,
        key_cache,
        value_cache,
        slot_mapping,
        kv_cache_dtype,
        k_scale,
        v_scale,
    )

# ...  Similarly, locate the reshape_and_cache function (non-Flash version) and make the same modifications  ...
def reshape_and_cache(
    key: torch.Tensor,
    value: torch.Tensor,
    key_cache: torch.Tensor,
    value_cache: torch.Tensor,
    slot_mapping: torch.Tensor,
    kv_cache_dtype: str,
    k_scale: torch.Tensor,
    v_scale: torch.Tensor,
) -> None:
    if kv_cache_dtype in ("fp8", "fp8_e4m3"):
        key = key.to(key_cache.dtype)
        value = value.to(value_cache.dtype)
        kv_cache_dtype = "auto"

    torch.ops._C_cache_ops.reshape_and_cache(
        key,
        value,
        key_cache,
        value_cache,
        slot_mapping,
        kv_cache_dtype,
        k_scale,
        v_scale,
    )

But I haven't tried it yet, so I don't know if it works

I successfully ran the program on 8*a100 40g. However, after installing VLLM using the code provided by the author, it is necessary to modify the “vllm/_custom_ops.py” file to enable support for kv_cache_dtype.

Thanks for your reply. Could you share your patch code? I would really appreciate it.

Find /usr/local/lib/python3.12/dist-packages/vllm/_custom_ops.py and modify “kv_cache_dtype” into “auto”,just like the answer above, finally "FP8" does not support. The LLM will work but it use huge GPU memory

Sign up or log in to comment