Skip to content

vllm.model_executor.kernels.linear.scaled_mm.humming

Classes:

HummingFP8ScaledMMLinearKernel

Bases: FP8ScaledMMLinearKernel

Humming GEMM Kernel for FP8.

Source code in vllm/model_executor/kernels/linear/scaled_mm/humming.py 
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
class HummingFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
    """Humming GEMM Kernel for FP8."""

    @classmethod
    def is_supported(
        cls, compute_capability: int | None = None
    ) -> tuple[bool, str | None]:
        if not current_platform.is_cuda():
            return False, "Humming only supported on CUDA"

        if not current_platform.has_device_capability(75):
            return False, "Humming only supported on SM75+"

        return True, None

    @classmethod
    def can_implement(
        cls, config: FP8ScaledMMLinearLayerConfig
    ) -> tuple[bool, str | None]:
        return True, None

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        from vllm.utils.humming import dtypes

        name_map = {"weight": "weight", "weight_scale": "weight_scale"}
        scale_torch_dtype = self.config.weight_quant_key.scale.dtype
        scale_dtype = dtypes.DataType.from_torch_dtype(scale_torch_dtype)

        quant_config = {
            "quant_method": "humming",
            "dtype": "float8e4m3",
            "scale_dtype": scale_dtype,
        }

        assert self.config.weight_quant_key.scale2 is None
        scale_group_shape = self.config.weight_quant_key.scale.group_shape
        if scale_group_shape.is_per_tensor():
            quant_config["weight_scale_type"] = "tensor"
            if not hasattr(layer, "global_scale") and hasattr(layer, "weight_scale"):
                del name_map["weight_scale"]
                name_map["global_scale"] = "weight_scale"
        elif scale_group_shape.is_per_channel():
            quant_config["weight_scale_type"] = "channel"
        elif scale_group_shape.is_per_group():
            quant_config["weight_scale_type"] = "group"
            quant_config["group_size"] = scale_group_shape.col
        else:
            assert scale_group_shape.row > 0 and scale_group_shape.col > 0
            quant_config["weight_scale_type"] = "block"
            quant_config["weight_scale_group_size_n"] = scale_group_shape.row
            quant_config["weight_scale_group_size"] = scale_group_shape.col

            if hasattr(layer, "weight_scale_inv"):
                name_map["weight_scale"] = "weight_scale_inv"

        convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
        prepare_humming_layer(layer, quant_config)

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        from vllm.utils.humming import HummingMethod

        flatten_inputs = x.view(-1, x.size(-1))
        output = HummingMethod.forward_layer(
            layer=layer,
            inputs=flatten_inputs,
            compute_config=layer.compute_config,
        )
        return output.view(*x.shape[:-1], output.size(-1))

    def apply_scaled_mm(
        self,
        *,
        A: torch.Tensor,
        B: torch.Tensor,
        out_dtype: torch.dtype,
        As: torch.Tensor,
        Bs: torch.Tensor,
        bias: torch.Tensor | None,
        output_shape: list,
    ) -> torch.Tensor:
        pass

HummingInt8ScaledMMLinearKernel

Bases: Int8ScaledMMLinearKernel

Humming GEMM Kernel for INT8.

Source code in vllm/model_executor/kernels/linear/scaled_mm/humming.py 
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
class HummingInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
    """Humming GEMM Kernel for INT8."""

    @classmethod
    def is_supported(
        cls, compute_capability: int | None = None
    ) -> tuple[bool, str | None]:
        if not current_platform.is_cuda():
            return False, "Humming only supported on CUDA"

        if not current_platform.has_device_capability(75):
            return False, "Humming only supported on SM75+"

        return True, None

    @classmethod
    def can_implement(
        cls, config: Int8ScaledMMLinearLayerConfig
    ) -> tuple[bool, str | None]:
        return True, None

    def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
        weight_name, weight_scale_name, *_ = self.layer_param_names
        name_map = {"weight": weight_name, "weight_scale": weight_scale_name}
        quant_config = {"quant_method": "humming", "dtype": "int8"}
        weight = getattr(layer, weight_name)
        weight.data = weight.data + 128

        convert_linear_layer_to_humming_standard(layer=layer, name_map=name_map)
        prepare_humming_layer(layer, quant_config)

    def apply_weights(
        self,
        layer: torch.nn.Module,
        x: torch.Tensor,
        bias: torch.Tensor | None = None,
    ) -> torch.Tensor:
        from vllm.utils.humming import HummingMethod

        flatten_inputs = x.view(-1, x.size(-1))
        output = HummingMethod.forward_layer(
            layer=layer,
            inputs=flatten_inputs,
            compute_config=layer.compute_config,
        )
        return output.view(*x.shape[:-1], output.size(-1))