[Bugfix] Get available quantization methods from quantization registry (#4098)

This commit is contained in:
Michael Goin
2024-04-18 03:21:55 -04:00
committed by GitHub
parent 66ded03067
commit 53b018edcb
6 changed files with 18 additions and 13 deletions

View File

@@ -16,13 +16,12 @@ from dataclasses import dataclass
import pytest
import torch
from vllm.model_executor.layers.quantization import (
_QUANTIZATION_CONFIG_REGISTRY)
from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS
capability = torch.cuda.get_device_capability()
capability = capability[0] * 10 + capability[1]
marlin_not_supported = (
capability < _QUANTIZATION_CONFIG_REGISTRY["marlin"].get_min_capability())
marlin_not_supported = (capability <
QUANTIZATION_METHODS["marlin"].get_min_capability())
@dataclass