A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 08:30 UTC
[Neuron] Support quantization on neuron (#18283)
Signed-off-by: Satyajith Chilappagari <satchill@amazon.com>
Satyajith Chilappagari committed
May 27, 2025 at 15:10 UTC
e0cbad4e30d7df4e9ee4634939ce19042639b735
3 files changed
+20
-2
tests/neuron/1_core/test_neuron_quant.py
new
+11
@@ -0,0 +1,11 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+from vllm.model_executor.layers.quantization.neuron_quant import (
3
+ NeuronQuantConfig)
4
+
5
+
6
+def test_get_supported_act_dtypes():
7
+ neuron_quant_config = NeuronQuantConfig()
8
+ supported_act_dtypes = neuron_quant_config.get_supported_act_dtypes()
9
+ target_list = ["any_dtype1", "any_dtype2"]
10
+ for dtype in target_list:
11
+ assert dtype in supported_act_dtypes
vllm/model_executor/layers/quantization/neuron_quant.py
+8
-1
@@ -13,6 +13,12 @@ from vllm.model_executor.layers.quantization.base_config import (
13
SUPPORTED_QUANT_DTYPE_LIST = ['s8', 'f8e4m3fn']
14
15
16
+class AlwaysSupportedDtypes(list):
17
+
18
+ def __contains__(self, item):
19
+ return True
20
+
21
+
22
class NeuronQuantConfig(QuantizationConfig):
23
"""Int8 Quantization Config class for Neuron Backend."""
24
@@ -35,7 +41,8 @@ class NeuronQuantConfig(QuantizationConfig):
41
return "neuron_quant"
42
43
def get_supported_act_dtypes(self) -> list[str]:
38
- return SUPPORTED_QUANT_DTYPE_LIST
44
+ # Neuron implements custom handling logic for quantization support
45
+ return AlwaysSupportedDtypes()
46
47
@classmethod
48
def get_min_capability(cls) -> int:
vllm/platforms/neuron.py
+1
-1
@@ -28,7 +28,7 @@ class NeuronPlatform(Platform):
28
device_name: str = "neuron"
29
device_type: str = "neuron"
30
ray_device_key: str = "neuron_cores"
31
- supported_quantization: list[str] = ["neuron_quant"]
31
+ supported_quantization: list[str] = ["neuron_quant", "fbgemm_fp8"]
32
device_control_env_var: str = "NEURON_RT_VISIBLE_CORES"
33
34
@classmethod