From 953973d5958b929b759c1ac960222a05d21d12d2 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Tue, 10 Sep 2024 14:51:01 +0000 Subject: [PATCH] env var for disabling kernels --- .../layers/quantization/kernels/__init__.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/kernels/__init__.py b/vllm/model_executor/layers/quantization/kernels/__init__.py index 22172771e5b64..88643ac512e00 100644 --- a/vllm/model_executor/layers/quantization/kernels/__init__.py +++ b/vllm/model_executor/layers/quantization/kernels/__init__.py @@ -1,3 +1,4 @@ +import os from typing import List, Optional, Type from vllm.platforms import current_platform @@ -24,6 +25,12 @@ def choose_mp_linear_kernel( failure_reasons = [] for kernel in _POSSIBLE_KERNELS: + if kernel.__name__ in os.environ.get("VLLM_DISABLED_KERNELS", "")\ + .split(","): + failure_reasons.append( + f' {kernel.__name__} disabled by environment variable') + continue + if kernel.get_min_capability() > compute_capability: failure_reasons.append( f"{kernel.__name__} requires capability " @@ -35,8 +42,7 @@ def choose_mp_linear_kernel( return kernel else: failure_reasons.append( - f' {kernel.__name__} cannot implement due to: {failure_reason}' - ) + f' {kernel.__name__} cannot implement due to: {failure_reason}') raise ValueError( "Failed to find a kernel that can implement the "\