From 6b26622f41b31763cc9a6be75e9231e691a18083 Mon Sep 17 00:00:00 2001 From: mgoin Date: Fri, 16 Feb 2024 10:23:36 -0500 Subject: [PATCH] Enable bfloat16 for sparse_w16a16 --- vllm/model_executor/layers/sparsity/sparse_w16a16.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/sparsity/sparse_w16a16.py b/vllm/model_executor/layers/sparsity/sparse_w16a16.py index d3a93d9b1d945..7a729ac2badd6 100644 --- a/vllm/model_executor/layers/sparsity/sparse_w16a16.py +++ b/vllm/model_executor/layers/sparsity/sparse_w16a16.py @@ -31,7 +31,7 @@ def get_name(cls) -> str: @classmethod def get_supported_act_dtypes(cls) -> List[torch.dtype]: - return [torch.half] + return [torch.float16, torch.bfloat16] @classmethod def get_min_capability(cls) -> int: