# Source code for tensorrt_llm.models.quantized.quant

    # SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
    # SPDX-License-Identifier: Apache-2.0
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    # http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    from typing import Any, Union

    import numpy as np

    from ...layers import ColumnLinear, RowLinear
    from ...models import (BaichuanForCausalLM, BloomForCausalLM, FalconForCausalLM,
                           GPTJForCausalLM, GPTLMHeadModel, LLaMAForCausalLM,
                           QWenForCausalLM)
    from ...module import Module
    from ...quantization import QuantMode
    from ...quantization.layers import FP8Linear, FP8RowLinear

    # isort: off
    from ...quantization.layers import (
        SmoothQuantAttention, SmoothQuantGatedMLP, SmoothQuantLayerNorm,
        SmoothQuantMLP, SmoothQuantRmsNorm, WeightOnlyGroupwiseQuantColumnLinear,
        WeightOnlyGroupwiseQuantRowLinear, WeightOnlyQuantColumnLinear,
        WeightOnlyQuantRowLinear)
    # isort: on

    def _smooth_quantize_gpt(model, quant_mode):
        assert quant_mode.has_act_and_weight_quant()
        for layer in model.layers:
            assert hasattr(layer,
                           "input_layernorm"), "The layer has no input_layernorm"
            layer.input_layernorm = SmoothQuantLayerNorm(
                normalized_shape=layer.hidden_size,
                dtype=layer.dtype,
                quant_mode=quant_mode)
            assert hasattr(layer, "attention"), "The layer has no attention"
            layer.attention = SmoothQuantAttention(
                layer.hidden_size,
                num_attention_heads=layer.num_attention_heads,
                max_position_embeddings=layer.max_position_embeddings,
                num_layers=layer.num_layers,
                apply_query_key_layer_scaling=layer.apply_query_key_layer_scaling,
                dtype=layer.dtype,
                attention_mask_type=layer.attention_mask_type,
                position_embedding_type=layer.position_embedding_type,
                tp_group=layer.tp_group,
                tp_size=layer.tp_size,
                quant_mode=quant_mode)
            assert hasattr(layer, "mlp"), "The layer has no mlp"
            layer.mlp = SmoothQuantMLP(hidden_size=layer.hidden_size,
                                       ffn_hidden_size=layer.hidden_size * 4,
                                       hidden_act=layer.hidden_act,
                                       dtype=layer.dtype,
                                       tp_group=layer.tp_group,
                                       tp_size=layer.tp_size,
                                       quant_mode=quant_mode)
            assert hasattr(layer,
                           "post_layernorm"), "The layer has no post_layernorm"
            layer.post_layernorm = SmoothQuantLayerNorm(
                normalized_shape=layer.hidden_size,
                dtype=layer.dtype,
                quant_mode=quant_mode)

        return model

    def _smooth_quantize_llama(model, quant_mode):
        assert quant_mode.has_act_and_weight_quant()
        for layer in model.layers:
            assert hasattr(layer,
                           "input_layernorm"), "The layer has no input_layernorm"
            layer.input_layernorm = SmoothQuantRmsNorm(
                normalized_shape=layer.hidden_size,
                dtype=layer.dtype,
                quant_mode=quant_mode)
            assert hasattr(layer, "attention"), "The layer has no attention"
            layer.attention = SmoothQuantAttention(
                layer.hidden_size,
                num_attention_heads=layer.num_attention_heads,
                num_kv_heads=layer.num_kv_heads,
                max_position_embeddings=layer.max_position_embeddings,
                num_layers=model.num_layers,
                dtype=layer.dtype,
                attention_mask_type=layer.attention_mask_type,
                position_embedding_type=layer.position_embedding_type,
                tp_group=layer.tp_group,
                tp_size=layer.tp_size,
                quant_mode=quant_mode,
                bias=False)

            assert hasattr(layer, "mlp"), "The layer has no mlp"
            layer.mlp = SmoothQuantGatedMLP(hidden_size=model.hidden_size,
                                            ffn_hidden_size=layer.mlp_hidden_size,
                                            hidden_act=layer.hidden_act,
                                            dtype=layer.dtype,
                                            tp_group=layer.tp_group,
                                            tp_size=layer.tp_size,
                                            quant_mode=quant_mode,
                                            bias=False)
            assert hasattr(
                layer,
                "post_layernorm"), "The layer has no post_rmspost_layernormnorm"
            layer.post_layernorm = SmoothQuantRmsNorm(
                normalized_shape=layer.hidden_size,
                dtype=layer.dtype,
                quant_mode=quant_mode)

        return model

    def _smooth_quantize_bloom(model, quant_mode):
        assert quant_mode.has_act_and_weight_quant()
        for layer in model.layers:
            assert hasattr(layer,
                           "input_layernorm"), "The layer has no input_layernorm"
            layer.input_layernorm = SmoothQuantLayerNorm(
                normalized_shape=layer.hidden_size,
                dtype=layer.dtype,
                quant_mode=quant_mode)
            assert hasattr(layer, "attention"), "The layer has no attention"
            layer.attention = SmoothQuantAttention(
                layer.hidden_size,
                num_attention_heads=layer.num_attention_heads,
                max_position_embeddings=layer.max_position_embeddings,
                num_layers=layer.num_layers,
                dtype=layer.dtype,
                attention_mask_type=layer.attention_mask_type,
                position_embedding_type=layer.position_embedding_type,
                tp_group=layer.tp_group,
                tp_size=layer.tp_size,
                tp_rank=layer.tp_rank,
                quant_mode=quant_mode)

            assert hasattr(layer, "mlp"), "The layer has no mlp"
            layer.mlp = SmoothQuantMLP(hidden_size=layer.hidden_size,
                                       ffn_hidden_size=layer.hidden_size * 4,
                                       hidden_act=layer.hidden_act,
                                       dtype=layer.dtype,
                                       tp_group=layer.tp_group,
                                       tp_size=layer.tp_size,
                                       quant_mode=quant_mode)
            assert hasattr(
                layer,
                "post_layernorm"), "The layer has no post_rmspost_layernormnorm"
            layer.post_layernorm = SmoothQuantLayerNorm(
                normalized_shape=layer.hidden_size,
                dtype=layer.dtype,
                quant_mode=quant_mode)

        setattr(model, 'quant_mode', quant_mode)
        return model

    def _smooth_quantize_baichuan(model, quant_mode):
        # Baichuan models' structures are similar to LLaMA so we can reuse the impl
        return _smooth_quantize_llama(model, quant_mode)

    def _smooth_quantize_internlm(model, quant_mode):
        assert quant_mode.has_act_and_weight_quant()
        for layer in model.layers:
            assert hasattr(layer,
                           "input_layernorm"), "The layer has no input_layernorm"
            layer.input_layernorm = SmoothQuantRmsNorm(
                normalized_shape=layer.hidden_size,
                dtype=layer.dtype,
                quant_mode=quant_mode)
            assert hasattr(layer, "attention"), "The layer has no attention"
            layer.attention = SmoothQuantAttention(
                layer.hidden_size,
                num_attention_heads=layer.num_attention_heads,
                num_kv_heads=layer.num_kv_heads,
                max_position_embeddings=layer.max_position_embeddings,
                num_layers=model.num_layers,
                dtype=layer.dtype,
                attention_mask_type=layer.attention_mask_type,
                position_embedding_type=layer.position_embedding_type,
                tp_group=layer.tp_group,
                tp_size=layer.tp_size,
                quant_mode=quant_mode,
                bias=model.attn_bias)

            assert hasattr(layer, "mlp"), "The layer has no mlp"
            layer.mlp = SmoothQuantGatedMLP(hidden_size=model.hidden_size,
                                            ffn_hidden_size=layer.mlp_hidden_size,
                                            hidden_act=layer.hidden_act,
                                            dtype=layer.dtype,
                                            tp_group=layer.tp_group,
                                            tp_size=layer.tp_size,
                                            quant_mode=quant_mode,
                                            bias=False)
            assert hasattr(
                layer,
                "post_layernorm"), "The layer has no post_rmspost_layernormnorm"
            layer.post_layernorm = SmoothQuantRmsNorm(
                normalized_shape=layer.hidden_size,
                dtype=layer.dtype,
                quant_mode=quant_mode)

        setattr(model, 'quant_mode', quant_mode)
        return model

    def _smooth_quantize_qwen(model, quant_mode):
        assert quant_mode.has_act_and_weight_quant()
        for layer in model.layers:
            assert hasattr(layer, "ln_1"), "The layer has no ln_1"
            layer.ln_1 = SmoothQuantRmsNorm(normalized_shape=layer.hidden_size,
                                            dtype=layer.dtype,
                                            quant_mode=quant_mode)
            assert hasattr(layer, "attention"), "The layer has no attention"
            layer.attention = SmoothQuantAttention(
                layer.hidden_size,
                layer.num_attention_heads,
                max_position_embeddings=layer.max_position_embeddings,
                num_layers=layer.num_layers,
                apply_query_key_layer_scaling=layer.apply_query_key_layer_scaling,
                attention_mask_type=layer.attention_mask_type,
                bias=layer.bias,
                qkv_bias_only=True,
                dtype=layer.dtype,
                position_embedding_type=layer.position_embedding_type,
                tp_group=layer.tp_group,
                tp_size=layer.tp_size,
                quant_mode=quant_mode)
            assert hasattr(layer, "mlp"), "The layer has no mlp"
            layer.mlp = SmoothQuantGatedMLP(hidden_size=layer.hidden_size,
                                            ffn_hidden_size=layer.mlp_hidden_size //
                                            2,
                                            hidden_act=layer.hidden_act,
                                            dtype=layer.dtype,
                                            bias=layer.bias,
                                            tp_group=layer.tp_group,
                                            tp_size=layer.tp_size,
                                            quant_mode=quant_mode)
            assert hasattr(layer, "ln_2"), "The layer has no ln_2"
            layer.ln_2 = SmoothQuantRmsNorm(normalized_shape=layer.hidden_size,
                                            dtype=layer.dtype,
                                            quant_mode=quant_mode)

        setattr(model, 'quant_mode', quant_mode)
        return model

    def _smooth_quantize(model, quant_mode):
        assert isinstance(model, GPTLMHeadModel) or isinstance(model, LLaMAForCausalLM) \
                or isinstance(model, BloomForCausalLM) or isinstance(model, BaichuanForCausalLM) \
                or isinstance(model, QWenForCausalLM), \
                "Only GPTLMHeadModel, LLaMAForCausalLM BloomForCausalLM and BaichuanForCausalLM are well tested now"
        if isinstance(model, GPTLMHeadModel):
            return _smooth_quantize_gpt(model, quant_mode)
        elif isinstance(model, LLaMAForCausalLM):
            return _smooth_quantize_llama(model, quant_mode)
        elif isinstance(model, BloomForCausalLM):
            return _smooth_quantize_bloom(model, quant_mode)
        elif isinstance(model, BaichuanForCausalLM):
            return _smooth_quantize_baichuan(model, quant_mode)
        elif isinstance(model, QWenForCausalLM):
            return _smooth_quantize_qwen(model, quant_mode)
        else:
            assert False, f"Model {type(model).__name__} is not supported by SmoothQuant yet"

    def _weight_only_quantize(model,
                              quant_mode,
                              exclude_modules=None,
                              current_key_name=None):
        assert quant_mode.is_weight_only()

        exclude_modules = ['lm_head'
                           ] if exclude_modules is None else exclude_modules

        for name, module in model.named_children():
            if current_key_name is None:
                current_key_name = []
            current_key_name.append(name)

            if len(list(module.children())) > 0:
                _weight_only_quantize(module, quant_mode, exclude_modules,
                                      current_key_name)

            if isinstance(module, ColumnLinear) and name not in exclude_modules:
                if not any(key in '.'.join(current_key_name)
                           for key in exclude_modules):
                    model._modules[name] = WeightOnlyQuantColumnLinear(
                        in_features=module.in_features,
                        out_features=module.out_features * module.tp_size,
                        bias=module.bias is not None,
                        dtype=module.dtype,
                        tp_group=module.tp_group,
                        tp_size=module.tp_size,
                        gather_output=module.gather_output,
                        quant_mode=quant_mode)
            elif isinstance(module, RowLinear) and name not in exclude_modules:
                if not any(key in '.'.join(current_key_name)
                           for key in exclude_modules):
                    model._modules[name] = WeightOnlyQuantRowLinear(
                        in_features=module.in_features * module.tp_size,
                        out_features=module.out_features,
                        bias=module.bias is not None,
                        dtype=module.dtype,
                        tp_group=module.tp_group,
                        tp_size=module.tp_size,
                        quant_mode=quant_mode)

            current_key_name.pop(-1)

        return model

    def _weight_only_groupwise_quantize(model,
                                        quant_mode,
                                        group_size=128,
                                        pre_quant_scale=False,
                                        zero=False,
                                        exclude_modules=None,
                                        current_key_name=None):
        exclude_modules = ['lm_head'
                           ] if exclude_modules is None else exclude_modules

        for name, module in model.named_children():
            if current_key_name is None:
                current_key_name = []
            current_key_name.append(name)

            if len(list(module.children())) > 0:
                _weight_only_groupwise_quantize(module, quant_mode, group_size,
                                                pre_quant_scale, zero,
                                                exclude_modules, current_key_name)

            if isinstance(module, ColumnLinear) and name not in exclude_modules:
                if not any(key in '.'.join(current_key_name)
                           for key in exclude_modules):
                    model._modules[name] = WeightOnlyGroupwiseQuantColumnLinear(
                        in_features=module.in_features,
                        out_features=module.out_features * module.tp_size,
                        group_size=group_size,
                        pre_quant_scale=pre_quant_scale,
                        zero=zero,
                        bias=module.bias is not None,
                        dtype=module.dtype,
                        tp_group=module.tp_group,
                        tp_size=module.tp_size,
                        gather_output=module.gather_output)
            elif isinstance(module, RowLinear) and name not in exclude_modules:
                if not any(key in '.'.join(current_key_name)
                           for key in exclude_modules):
                    model._modules[name] = WeightOnlyGroupwiseQuantRowLinear(
                        in_features=module.in_features * module.tp_size,
                        out_features=module.out_features,
                        group_size=group_size,
                        pre_quant_scale=pre_quant_scale,
                        zero=zero,
                        bias=module.bias is not None,
                        dtype=module.dtype,
                        tp_group=module.tp_group,
                        tp_size=module.tp_size)

            current_key_name.pop(-1)

        return model

    def quantize_model(model: Module, quant_mode: QuantMode, **kwargs: Any):
        if quant_mode.has_fp8_qdq() or quant_mode.has_fp8_kv_cache():
            model = _fp8_quantize(model, quant_mode, **kwargs)
        elif quant_mode.has_act_and_weight_quant():
            model = _smooth_quantize(model, quant_mode)
        elif quant_mode.is_weight_only():
            if quant_mode.has_per_group_scaling():
                model = _weight_only_groupwise_quantize(model, quant_mode, **kwargs)
            else:
                model = _weight_only_quantize(model, quant_mode, **kwargs)

        setattr(model, "quant_mode", quant_mode)
        return model

    def get_dummy_quant_scales(num_layers):
        return {
            'lm_head_act': 0.99,
            'lm_head_weights': 0.99,
            'fc_act': [0.99 for _ in range(num_layers)],
            'fc_weights': [0.99 for _ in range(num_layers)],
            'gate_act': [0.99 for _ in range(num_layers)],
            'gate_weights': [0.99 for _ in range(num_layers)],
            'proj_act': [0.99 for _ in range(num_layers)],
            'proj_weights': [0.99 for _ in range(num_layers)],
            'qkv_act': [0.99 for _ in range(num_layers)],
            'qkv_weights': [0.99 for _ in range(num_layers)],
            'qkv_output': [5.0 for _ in range(num_layers)],
            'dense_act': [0.99 for _ in range(num_layers)],
            'dense_weights': [0.99 for _ in range(num_layers)],
        }

    def _quantize_layer(layer, layer_idx, quant_mode, quant_scales):
        assert hasattr(layer, "mlp"), "The layer has no mlp"
        fake_fp8_sf_dt = np.float32

        assert isinstance(layer.mlp.fc, (FP8Linear, FP8RowLinear))
        assert isinstance(layer.mlp.proj, (FP8Linear, FP8RowLinear))
        layer.mlp.fc.activation_scaling_factor.value = np.array(
            [quant_scales['fc_act'][layer_idx]], dtype=fake_fp8_sf_dt)
        layer.mlp.fc.weights_scaling_factor.value = np.array(
            [quant_scales['fc_weights'][layer_idx]], dtype=fake_fp8_sf_dt)
        layer.mlp.proj.activation_scaling_factor.value = np.array(
            [quant_scales['proj_act'][layer_idx]], dtype=fake_fp8_sf_dt)
        layer.mlp.proj.weights_scaling_factor.value = np.array(
            [quant_scales['proj_weights'][layer_idx]], dtype=fake_fp8_sf_dt)
        if hasattr(layer.mlp, 'gate'):
            assert isinstance(layer.mlp.gate, (FP8Linear, FP8RowLinear))
            layer.mlp.gate.activation_scaling_factor.value = np.array(
                [quant_scales['gate_act'][layer_idx]], dtype=fake_fp8_sf_dt)
            layer.mlp.gate.weights_scaling_factor.value = np.array(
                [quant_scales['gate_weights'][layer_idx]], dtype=fake_fp8_sf_dt)

        assert hasattr(layer, "attention"), "The layer has no attention"
        assert isinstance(layer.attention.qkv, (FP8Linear, FP8RowLinear))
        assert isinstance(layer.attention.dense, (FP8Linear, FP8RowLinear))
        layer.attention.qkv.activation_scaling_factor.value = np.array(
            [quant_scales['qkv_act'][layer_idx]], dtype=fake_fp8_sf_dt)
        layer.attention.qkv.weights_scaling_factor.value = np.array(
            [quant_scales['qkv_weights'][layer_idx]], dtype=fake_fp8_sf_dt)
        if quant_mode.has_fp8_kv_cache():
            layer.attention.kv_orig_quant_scale.value = np.array(
                [quant_scales['qkv_output'][layer_idx]], dtype=fake_fp8_sf_dt)
            layer.attention.kv_quant_orig_scale.value = np.array(
                [1.0 / quant_scales['qkv_output'][layer_idx]], dtype=fake_fp8_sf_dt)
        layer.attention.dense.activation_scaling_factor.value = np.array(
            [quant_scales['dense_act'][layer_idx]], dtype=fake_fp8_sf_dt)
        layer.attention.dense.weights_scaling_factor.value = np.array(
            [quant_scales['dense_weights'][layer_idx]], dtype=fake_fp8_sf_dt)

        return layer

    def _default_fp8_quantize(model: Union[GPTLMHeadModel, LLaMAForCausalLM,
                                           GPTJForCausalLM],
                              quant_mode: QuantMode,
                              quant_scales: dict = None):
        """
        Quantize all linear layers (i.e., MLP, Attention QKV/Dense) and KV cache IO with dummy scales
        This is used by benchmark script and therefore is intentionally decoupled from AMMO toolkit
        """
        if quant_scales is None:
            num_layers = getattr(model, '_num_layers',
                                 getattr(model, 'num_layers', None))
            assert num_layers is not None
            quant_scales = get_dummy_quant_scales(num_layers)

        assert model.quant_mode == quant_mode, "Quant setting not consistent with model init setting"

        use_fp8_qdq = quant_mode.has_fp8_qdq()
        assert use_fp8_qdq

        for layer_idx, layer in enumerate(model.layers):
            layer = _quantize_layer(layer, layer_idx, quant_mode, quant_scales)

        # TODO: add lm_head

        return model

    def _fp8_quantize(model, quant_mode: QuantMode, quant_scales: dict = None):
        if isinstance(
                model,
            (FalconForCausalLM, GPTJForCausalLM, GPTLMHeadModel, LLaMAForCausalLM)):
            return _default_fp8_quantize(model, quant_mode, quant_scales)
        raise NotImplementedError(
            f"Model {model} is not implemented by fp8_quantize yet")