From 3c11f56db6bc18558cd3b738e2f25e231a1d703a Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Fri, 16 Feb 2024 07:28:02 -0800 Subject: [PATCH] Enable bfloat16 for sparse_w16a16 (#18) --- 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: