per token per channel quantization initial setup

Author: kliuae

vllm/config.py

@@ -427,7 +427,7 @@ def _verify_quantization(self) -> None:
         optimized_quantization_methods = [
             "fp8", "marlin", "modelopt", "gptq_marlin_24", "gptq_marlin",
             "awq_marlin", "fbgemm_fp8", "compressed_tensors",
-            "compressed-tensors", "experts_int8"
+            "compressed-tensors", "experts_int8", "ptpc_fp8"
         ]
         if self.quantization is not None:
             self.quantization = self.quantization.lower()

vllm/model_executor/layers/quantization/__init__.py

@@ -10,6 +10,7 @@
     "tpu_int8",
     "fp8",
     "fbgemm_fp8",
+    "ptpc_fp8",
     "modelopt",
     # The order of gptq methods is important for config.py iteration over
     # override_quantization_method(..)
@@ -43,6 +44,7 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
     from .deepspeedfp import DeepSpeedFPConfig
     from .experts_int8 import ExpertsInt8Config
     from .fbgemm_fp8 import FBGEMMFp8Config
+    from .ptpc_fp8 import PTPCFp8Config
     from .fp8 import Fp8Config
     from .gguf import GGUFConfig
     from .gptq import GPTQConfig
@@ -63,6 +65,7 @@ def get_quantization_config(quantization: str) -> Type[QuantizationConfig]:
         "tpu_int8": Int8TpuConfig,
         "fp8": Fp8Config,
         "fbgemm_fp8": FBGEMMFp8Config,
+        "ptpc_fp8": PTPCFp8Config,
         "modelopt": ModelOptFp8Config,
         # The order of gptq methods is important for config.py iteration over
         # override_quantization_method(..)

vllm/model_executor/layers/quantization/ptpc_fp8.py

@@ -0,0 +1,50 @@
+from typing import Any, Dict, List, Optional
+
+import torch
+from torch.nn import Module
+from torch.nn.parameter import Parameter
+
+from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
+                                               UnquantizedLinearMethod)
+from vllm.model_executor.layers.quantization.base_config import (
+    QuantizationConfig, QuantizeMethodBase)
+from vllm.model_executor.layers.quantization.fp8 import cutlass_fp8_supported
+from vllm.model_executor.layers.quantization.fbgemm_fp8 import (
+    FBGEMMFp8Config, FBGEMMFp8LinearMethod)
+from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
+    apply_fp8_marlin_linear, prepare_fp8_layer_for_marlin)
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
+    is_layer_skipped)
+from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
+    apply_fp8_linear, normalize_e4m3fn_to_e4m3fnuz)
+from vllm.model_executor.parameter import (ChannelQuantScaleParameter,
+                                           ModelWeightParameter)
+from vllm.platforms import current_platform
+
+logger = init_logger(__name__)
+
+
+class PTPCFp8Config(FBGEMMFp8Config):
+    """Config class for Per-Token-Per-Channel Fp8."""
+
+    def __init__(self, ignore_list: Optional[List[str]] = None):
+        super().__init__(ignore_list, 1.0) # Dummy values
+
+    @classmethod
+    def get_name(cls) -> str:
+        return "ptpc_fp8"
+
+    @classmethod
+    def from_config(cls, config: Dict[str, Any]) -> "PTPCFp8Config":
+        ignore_list = cls.get_from_keys(config, ["modules_to_not_convert"])
+        input_scale_ub = cls.get_from_keys(config, ["activation_scale_ub"])
+        return cls(ignore_list=ignore_list, input_scale_ub=input_scale_ub)
+
+    def get_quant_method(self, layer: torch.nn.Module,
+                         prefix: str) -> Optional["QuantizeMethodBase"]:
+        if isinstance(layer, LinearBase):
+            if is_layer_skipped(prefix, self.ignore_list):
+                return UnquantizedLinearMethod()
+            return FBGEMMFp8LinearMethod(self)
+        return None

vllm/platforms/rocm.py

@@ -79,7 +79,7 @@ class RocmPlatform(Platform):
     dispatch_key: str = "CUDA"
     supported_quantization: list[str] = [
         "awq", "gptq", "fp8", "compressed_tensors", "compressed-tensors",
-        "fbgemm_fp8", "gguf"
+        "fbgemm_fp8", "gguf", "ptpc_fp8"
     ]
 
     @classmethod