Consolidate quant method names into a single file (#1101)

Author: kyuyeunk

tpu_inference/layers/common/quant_methods.py

@@ -0,0 +1,8 @@
+UNQUANTIZED = "unquantized"
+MXFP4 = "mxfp4"
+AWQ = "awq"
+COMPRESSED_TENSORS = "compressed-tensors"
+
+
+def get_tpu_quant_method(quant_method: str) -> str:
+    return "tpu-" + quant_method

tpu_inference/layers/vllm/quantization/__init__.py

@@ -5,6 +5,7 @@
 from vllm.model_executor.layers.quantization.base_config import \
     QuantizationConfig
 
+from tpu_inference.layers.common import quant_methods
 from tpu_inference.layers.vllm.quantization.awq import VllmAWQConfig
 from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
 from tpu_inference.layers.vllm.quantization.compressed_tensors.compressed_tensors import \
@@ -20,9 +21,9 @@ def get_tpu_quantization_config(vllm_config: VllmConfig,
     # TODO(kyuyeunk): Add support for "tpu_int8".
     method_to_config: dict[str, str] = {
         None: VllmUnquantizedConfig,
-        "compressed-tensors": VllmCompressedTensorsConfig,
-        "awq": VllmAWQConfig,
-        "mxfp4": VllmMxfp4Config,
+        quant_methods.COMPRESSED_TENSORS: VllmCompressedTensorsConfig,
+        quant_methods.AWQ: VllmAWQConfig,
+        quant_methods.MXFP4: VllmMxfp4Config,
     }
     if model_config.quantization not in method_to_config:
         raise NotImplementedError(
@@ -32,7 +33,7 @@ def get_tpu_quantization_config(vllm_config: VllmConfig,
     assert issubclass(quant_config, JaxCommonConfig)
     quant_config.set_configs(vllm_config, mesh)
 
-    # TODO(kyuyeunk): Create more programmatic way to handle this.
-    model_config.quantization = "tpu-" + quant_config.get_name()
+    model_config.quantization = quant_methods.get_tpu_quant_method(
+        quant_config.get_name())
     return VllmConfig.get_quantization_config(model_config,
                                               vllm_config.load_config)

tpu_inference/layers/vllm/quantization/awq.py

@@ -18,6 +18,7 @@
     is_layer_skipped, unpack_quantized_values_into_int32)
 from vllm.scalar_type import scalar_types
 
+from tpu_inference.layers.common.quant_methods import AWQ, get_tpu_quant_method
 from tpu_inference.layers.vllm.linear_common import (
     slice_sharded_tensor_for_concatenation, torch_to_jax_param)
 from tpu_inference.layers.vllm.quantization.common import (
@@ -29,9 +30,13 @@
 logger = init_logger(__name__)
 
 
-@register_quantization_config("tpu-awq")
+@register_quantization_config(get_tpu_quant_method(AWQ))
 class VllmAWQConfig(AWQConfig, JaxCommonConfig):
 
+    @classmethod
+    def get_name(cls):
+        return AWQ
+
     def get_supported_act_dtypes(self) -> list[torch.dtype]:
         # NOTE: AWQ checkpoint was quantized with float16. But on TPUs, using
         # bfloat16 is signifcantly preferred over foat16. This might lead to

tpu_inference/layers/vllm/quantization/compressed_tensors/compressed_tensors.py

@@ -16,6 +16,8 @@
 from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
     find_matched_target, should_ignore_layer)
 
+from tpu_inference.layers.common.quant_methods import (COMPRESSED_TENSORS,
+                                                       get_tpu_quant_method)
 from tpu_inference.layers.vllm.quantization.common import JaxCommonConfig
 from tpu_inference.layers.vllm.quantization.compressed_tensors.compressed_tensors_moe import \
     VllmCompressedTensorsW8A8Fp8MoEMethod
@@ -30,9 +32,13 @@
 logger = init_logger(__name__)
 
 
-@register_quantization_config("tpu-compressed-tensors")
+@register_quantization_config(get_tpu_quant_method(COMPRESSED_TENSORS))
 class VllmCompressedTensorsConfig(CompressedTensorsConfig, JaxCommonConfig):
 
+    @classmethod
+    def get_name(cls) -> str:
+        return COMPRESSED_TENSORS
+
     def get_scheme(self,
                    layer: torch.nn.Module,
                    layer_name: Optional[str] = None

tpu_inference/layers/vllm/quantization/mxfp4.py

@@ -24,6 +24,8 @@
 from vllm.model_executor.layers.quantization.utils.quant_utils import \
     is_layer_skipped
 
+from tpu_inference.layers.common.quant_methods import (MXFP4,
+                                                       get_tpu_quant_method)
 from tpu_inference.layers.vllm.fused_moe import fused_moe_func_padded
 from tpu_inference.layers.vllm.linear_common import \
     reorder_concatenated_tensor_for_sharding
@@ -64,9 +66,13 @@ def dequantize_block_weight(weight: jax.Array,
     return weight_dequantized.reshape(orig_shape).astype(out_dtype)
 
 
-@register_quantization_config("tpu-mxfp4")
+@register_quantization_config(get_tpu_quant_method(MXFP4))
 class VllmMxfp4Config(Mxfp4Config, JaxCommonConfig):
 
+    @classmethod
+    def get_name(cls):
+        return MXFP4
+
     def get_quant_method(self, layer: torch.nn.Module,
                          prefix: str) -> Optional["QuantizeMethodBase"]:
         from vllm.attention.layer import Attention  # Avoid circular import

tpu_inference/layers/vllm/quantization/unquantized.py

@@ -23,6 +23,8 @@
 
 from tpu_inference import envs
 from tpu_inference.kernels.fused_moe.v1.kernel import fused_ep_moe
+from tpu_inference.layers.common.quant_methods import (UNQUANTIZED,
+                                                       get_tpu_quant_method)
 from tpu_inference.layers.vllm.fused_moe import fused_moe_func_padded
 from tpu_inference.layers.vllm.linear_common import (
     reorder_concatenated_tensor_for_sharding,
@@ -34,12 +36,12 @@
 logger = init_logger(__name__)
 
 
-@register_quantization_config("tpu-unquantized")
+@register_quantization_config(get_tpu_quant_method(UNQUANTIZED))
 class VllmUnquantizedConfig(QuantizationConfig, JaxCommonConfig):
 
     @classmethod
     def get_name(cls) -> str:
-        return "unquantized"
+        return UNQUANTIZED
 
     @classmethod
     def get_supported_act_dtypes(cls) -> list[torch.dtype]: