builder.py

# SPDX-FileCopyrightText: Copyright (c) 2022-2024 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.
import copy
import json
import math
import os
import shutil
import time
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, Optional, Union

import numpy as np
import tensorrt as trt

from ._common import _is_building, check_max_num_tokens, serialize_engine
from ._utils import np_bfloat16, np_float8, str_dtype_to_trt, to_json_file
from .auto_parallel import auto_parallel
from .auto_parallel.config import AutoParallelConfig
from .bindings import KVCacheType
from .functional import PositionEmbeddingType
from .graph_rewriting import optimize
from .logger import logger
from .lora_manager import LoraConfig
from .models import PretrainedConfig, PretrainedModel
from .models.modeling_utils import SpeculativeDecodingMode, optimize_model
from .network import Network, net_guard
from .plugin import PluginConfig
from .quantization import QuantAlgo, QuantMode
from .version import __version__


class ConfigEncoder(json.JSONEncoder):

    def default(self, obj):
        if isinstance(obj, KVCacheType):
            # For KVCacheType, convert it to string by split of 'KVCacheType.PAGED'.
            return obj.__str__().split('.')[-1]
        else:
            return super().default(obj)


class BuilderConfig(object):

    def __init__(self, **kwargs):
        # intentionally use **kwargs, user should never call this ctor directly,
        # use Builder.create_builder_config() instead
        pass

    def _init(self, trt_builder_config, **kwargs):
        self._trt_builder_config = trt_builder_config
        for key, value in kwargs.items():
            setattr(self, key, value)
        return self

    @property
    def trt_builder_config(self) -> trt.IBuilderConfig:
        return self._trt_builder_config

    def to_dict(self) -> Dict:
        '''return a dict with keys
        {
            "builder_config": {
                # all key values set by the _init function
            },
            "plugin_config": {
                # the network plugin_config (if any) attached to this BuilderConfig object
                # inside the Builder.build_engine
            },
            "auto_parallel_config": {
                # the network auto_parallel_config (if any) attached to this BuilderConfig object
                # inside the Builder.build_engine
            }
        }
        '''
        config = {'builder_config': {}}
        for k in self.__dict__.keys():
            if k not in [
                    '_trt_builder_config', 'plugin_config',
                    'auto_parallel_config'
            ]:
                config['builder_config'][k] = self.__getattribute__(k)
        if hasattr(self, 'plugin_config'):
            assert isinstance(self.plugin_config, PluginConfig), \
                f"Found unexpected plugin_config object with type: {type(self.plugin_config)}"
            config['plugin_config'] = self.plugin_config.to_dict()
        return config


class Builder():

    _ALLOWED_PRECISIONS = [
        'float32', 'float16', 'bfloat16', trt.DataType.HALF, trt.DataType.FLOAT,
        trt.DataType.BF16
    ]

    def __init__(self):
        super().__init__()
        self._trt_builder = trt.Builder(logger.trt_logger)
        # TODO: Enable strongly_typed on by default in TRT 10.0
        self.strongly_typed = False

    @property
    def trt_builder(self) -> trt.Builder:
        return self._trt_builder

    def create_network(self) -> Network:
        explicit_batch_flag = 0
        # Explicit batch flag will be deprecated in TRT 10
        if "EXPLICIT_BATCH" in trt.NetworkDefinitionCreationFlag.__members__.keys(
        ):
            explicit_batch_flag = 1 << int(
                trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)

        if self.strongly_typed:
            return Network()._init(
                self.trt_builder.create_network(
                    explicit_batch_flag
                    | (1 << int(
                        trt.NetworkDefinitionCreationFlag.STRONGLY_TYPED))))
        else:
            return Network()._init(
                self.trt_builder.create_network(explicit_batch_flag))

    def create_builder_config(self,
                              precision: Union[str, trt.DataType],
                              timing_cache: Union[str, Path,
                                                  trt.ITimingCache] = None,
                              tensor_parallel: int = 1,
                              use_refit: bool = False,
                              int8: bool = False,
                              strongly_typed: bool = False,
                              opt_level: Optional[int] = None,
                              force_num_profiles: Optional[int] = None,
                              profiling_verbosity: str = "layer_names_only",
                              use_strip_plan: bool = False,
                              weight_streaming: bool = False,
                              **kwargs) -> BuilderConfig:
        ''' @brief Create a builder config with given precisions and timing cache
            @param precision: one of allowed precisions, defined in Builder._ALLOWED_PRECISIONS
            @param timing_cache: a timing cache object or a path to a timing cache file
            @param tensor_parallel: number of GPUs used for tensor parallel
            @param kwargs: any other arguments users would like to attach to the config object as attributes
            @param refit: set to accelerate multi-gpu building, build engine for 1 gpu and refit for the others
            @param int8: whether to build with int8 enabled or not. Can't be used together with refit option
            @return: A BuilderConfig object, return None if failed
        '''
        self.strongly_typed = self.strongly_typed or strongly_typed

        quant_mode = kwargs.get("quant_mode", QuantMode(0))
        if not strongly_typed and precision not in self._ALLOWED_PRECISIONS:
            logger.error(
                f"precision should be one of {self._ALLOWED_PRECISIONS}")

        config = self.trt_builder.create_builder_config()
        if weight_streaming:
            config.set_flag(trt.BuilderFlag.WEIGHT_STREAMING)
        if not self.strongly_typed:
            fp8 = quant_mode.has_fp8_qdq() or quant_mode.has_fp8_kv_cache()
            if precision == 'float16' or precision == trt.DataType.HALF:
                config.set_flag(trt.BuilderFlag.FP16)
                config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
            elif precision == 'bfloat16' or precision == trt.DataType.BF16:
                config.set_flag(trt.BuilderFlag.BF16)
                config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)
            if int8:
                config.set_flag(trt.BuilderFlag.INT8)

            if fp8:
                config.set_flag(trt.BuilderFlag.FP8)
                config.set_flag(trt.BuilderFlag.OBEY_PRECISION_CONSTRAINTS)

        config.set_preview_feature(trt.PreviewFeature.PROFILE_SHARING_0806,
                                   True)

        if use_refit:
            config.set_flag(trt.BuilderFlag.REFIT)

        # Use fine-grained refit when strip plan is enabled in TRT10.2+.
        if use_strip_plan:
            config.set_flag(trt.BuilderFlag.REFIT_INDIVIDUAL)

        if use_strip_plan:
            config.set_flag(trt.BuilderFlag.STRIP_PLAN)

        if opt_level is not None:
            config.builder_optimization_level = opt_level

        # Set TRT Engine profiling verbosity
        if profiling_verbosity == "detailed":
            config.profiling_verbosity = trt.ProfilingVerbosity.DETAILED
        elif profiling_verbosity == "none":
            config.profiling_verbosity = trt.ProfilingVerbosity.NONE
        else:
            config.profiling_verbosity = trt.ProfilingVerbosity.LAYER_NAMES_ONLY

        # set timing cache
        cache = None
        if timing_cache is not None:
            # use given cache
            if isinstance(timing_cache, trt.ITimingCache):
                cache = timing_cache
            # read cache from file
            elif isinstance(timing_cache,
                            (str, Path)) and os.path.exists(timing_cache):
                with open(timing_cache, "rb") as f:
                    cache = config.create_timing_cache(f.read())
            else:
                logger.warning(
                    "Invalid timing cache, using freshly created one")
        if cache is None:
            cache = config.create_timing_cache(b"")
        # When user does not given any existing cache, internally always created one
        # so the cache should never None here
        assert cache is not None and isinstance(cache, trt.ITimingCache)
        config.set_timing_cache(cache, ignore_mismatch=False)

        # set weight sparsity
        weight_sparsity = kwargs.get("weight_sparsity", False)
        if weight_sparsity:
            config.set_flag(trt.BuilderFlag.SPARSE_WEIGHTS)

        return BuilderConfig()._init(config,
                                     precision=precision,
                                     tensor_parallel=tensor_parallel,
                                     use_refit=use_refit,
                                     int8=int8,
                                     force_num_profiles=force_num_profiles,
                                     strongly_typed=self.strongly_typed,
                                     use_strip_plan=use_strip_plan,
                                     **kwargs)

    def _add_optimization_profile(self, network: Network,
                                  builder_config: BuilderConfig):
        assert isinstance(builder_config, BuilderConfig)
        assert isinstance(network, Network)
        input_tensors = network._inputs
        if len(input_tensors) == 0:
            logger.warning("There are no inputs in the network!")
            return
        num_profiles = len(list(input_tensors.values())[0].profiles)
        force_num_profiles = getattr(
            builder_config, "force_num_profiles") if hasattr(
                builder_config, "force_num_profiles") else None
        for i in range(num_profiles):
            logger.debug(f'Adding optimization profile {i+1}/{num_profiles}')
            profile = self.trt_builder.create_optimization_profile()
            for input_name in input_tensors.keys():
                if len(input_tensors[input_name].profiles) == 0:
                    continue
                shape_profile = input_tensors[input_name].profiles[i]
                min_shape = [*shape_profile.min]
                opt_shape = [*shape_profile.opt]
                max_shape = [*shape_profile.max]
                if network._auto_parallel_config is not None:
                    io_shards = network._auto_parallel_config["io_shards"]
                    if input_name in io_shards:
                        shards = io_shards[input_name]
                        for dim, shard_num in shards.items():
                            min_shape[dim] = int(
                                math.floor(min_shape[dim] / shard_num))
                            opt_shape[dim] = int(
                                round(opt_shape[dim] / shard_num))
                            max_shape[dim] = int(
                                math.ceil(max_shape[dim] / shard_num))
                profile.set_shape(input_name, min_shape, opt_shape, max_shape)
                logger.debug(
                    f'{input_name}, min: {min_shape}, opt: {opt_shape}, max: {max_shape}, dimension names: {shape_profile.dimension_names}'
                )
            ret = builder_config.trt_builder_config.add_optimization_profile(
                profile)
            logger.debug(f"Added optimization profile: #{ret}")
            if force_num_profiles is not None and (
                    i + 1
            ) == force_num_profiles and force_num_profiles < num_profiles:
                logger.warning(
                    f"Only adding {force_num_profiles} profiles instead of {num_profiles}."
                )
                break
        assert self._validate_named_dimensions(
            network, builder_config
        ), "Validation of the tensor dimension ranges failed, please check the dimension ranges, find the offensive tensor and dimension name in above the error log"

    def _validate_named_dimensions(self, network: Network,
                                   builder_config) -> bool:
        '''
            For each profile, validate that the named dimensions of different input tensors in this profile all have same range.
            TRT will validate the same condition, validate it earlier to make sure the modeling in TensorRT-LLM are correct and
            makes the error msg more user friendly.
        '''
        valid = True
        for profile_idx in range(
                builder_config.trt_builder_config.num_optimization_profiles):
            dimension_to_range = {}
            for input_name, input_tensor in network._inputs.items():
                # it's legal that a Tensor does not have dim_range?
                if len(input_tensor.profiles) != 0:
                    profile = input_tensor.profiles[profile_idx]
                    for dim_idx, dim_name in enumerate(profile.dimension_names):
                        if dim_name not in dimension_to_range:
                            dimension_to_range[dim_name] = []
                        min, opt, max = profile.min[dim_idx], profile.opt[
                            dim_idx], profile.max[dim_idx]
                        dimension_to_range[dim_name].append(
                            (input_name, (min, opt, max)))
            for dim, ranges in dimension_to_range.items():
                unique_ranges = set([r[1] for r in ranges])
                logger.debug(
                    f"Validating dimension:{dim}, ranges for this dim are:{unique_ranges}"
                )
                if len(unique_ranges) != 1:
                    logger.error(
                        f"Found illegal dimension setting for profile {profile_idx}, dimension name is: {dim}"
                    )
                    logger.error(
                        "Offensive tensors which have this dimension are:\n" +
                        "\n".join([f"{r[1]} {dim} {r[0]}" for r in ranges]))
                    valid = False
        return valid

    @_is_building
    def refit_engine(self, network: Network, engine_buffer) -> trt.IHostMemory:
        '''
            @brief: Refit one TensorRT engine using weights from the network,
                user should guarantee that the engine is built with REFIT flag, and the network has the same structure with the engine.
            @param engine_buffer: A serialized TensorRT engine.
            @param network: Network object.
            @return: A serialized TRT engine if refit successfully, None otherwise
        '''
        assert isinstance(network, Network)
        logger.info('Refit TRT engine')
        runtime = trt.Runtime(logger.trt_logger)
        engine = runtime.deserialize_cuda_engine(engine_buffer)

        tik = time.time()

        # Refit engine
        refitter = trt.Refitter(engine, logger.trt_logger)
        if network.named_parameters is not None:
            for name, param in network.named_parameters:
                if param._get_weights(
                ) is None or not refitter.set_named_weights(
                        name, param._get_weights()):
                    logger.error(f'Failed to refit weight: {name}')
                    return None
        else:
            logger.error(
                'Please set named parameters before building multiple engines.')
            return None

        if not refitter.refit_cuda_engine():
            logger.error('Failed to refit engine.')
            return None

        tok = time.time()
        t = time.strftime('%H:%M:%S', time.gmtime(tok - tik))
        logger.info(f'Total time of refitting {engine.name}: {t}')
        serialized_engine = engine.serialize()
        return serialized_engine

    @_is_building
    def build_engine(self,
                     network: Network,
                     builder_config: BuilderConfig,
                     managed_weights: list = None) -> trt.IHostMemory:
        '''
            @brief: Build one TensorRT engine from the network.
            @param network: Network object.
            @param builder_config: BuilderConfig object.
            @return: A serialized TRT engine.
        '''
        assert isinstance(network, Network)
        builder_config.plugin_config = network.plugin_config
        builder_config.auto_parallel_config = network.auto_parallel_config
        if builder_config.trt_builder_config.num_optimization_profiles == 0:
            self._add_optimization_profile(network, builder_config)
        logger.info(
            f"Total optimization profiles added: {builder_config.trt_builder_config.num_optimization_profiles}"
        )
        engine = None

        # Rename weights
        if network.named_parameters is not None:
            managed_parameters = []
            for name, param in network.named_parameters:
                if param.is_managed(network):
                    assert managed_weights is not None, "managed_weights should be provided when enabled"
                    managed_parameters.append(param)
                    param.set_name(name, network)
                    continue
                if param._get_weights(network) is None:
                    if not param.is_buffer:
                        logger.info(
                            f"Parameter {name} {param.raw_value.shape} {param.raw_value.dtype} was created"
                            " but unused in forward method, so not materialized to TRT network"
                        )
                    continue
                if not param.set_name(name, network):
                    raise RuntimeError(f'Failed to set weight: {name}')
                # This mark_weights_refittable has no side effect when refit_individual is not enabled.
                network.trt_network.mark_weights_refittable(name)

        network._fill_weights()

        # Build engine
        logger.info(f'Build TensorRT engine {network.trt_network.name}')
        tik = time.time()
        engine = self.trt_builder.build_serialized_network(
            network.trt_network, builder_config.trt_builder_config)
        assert engine is not None, 'Engine building failed, please check the error log.'

        tok = time.time()
        t = time.strftime('%H:%M:%S', time.gmtime(tok - tik))
        logger.info(f'Total time of building {network.trt_network.name}: {t}')

        if managed_weights is not None and network.named_parameters is not None:
            for param in managed_parameters:
                name = param.name
                value: np.ndarray = param._value
                if value is None:
                    logger.error(f'Failed to get weight: {name}')
                    continue
                if value.dtype == np.float16 and value.ndim == 2 and network.plugin_config.gemm_plugin is None:
                    # MOE has ndim=3 and uses plugin, no need to transpose
                    value = value.transpose(1, 0)  # WAR for bug 4641821
                managed_weights.append((name, value))

        return engine

    @staticmethod
    def save_timing_cache(builder_config: BuilderConfig, out_path: str) -> bool:
        '''Serialize timing cache of given builder config to file specified by out_path
            return True if the cache is successfully serialized, False otherwise
        '''
        cache = builder_config.trt_builder_config.get_timing_cache()
        if cache is None:
            logger.warning(
                'No timing cache found in the given builder config, skip saving.'
            )
            return False
        with cache.serialize() as buffer:
            with open(out_path, "wb") as f:
                f.write(buffer)
                f.flush()
                os.fsync(f)
        logger.info(f'Timing cache serialized to {out_path}')
        return True

    @staticmethod
    def save_config(builder_config: BuilderConfig, config_path: str):
        config = builder_config.to_dict()
        to_json_file(config, config_path)
        logger.info(f'Config saved to {config_path}.')


@dataclass
class BuildConfig:
    max_input_len: int = 256
    max_seq_len: int = 512
    opt_batch_size: int = 8
    max_batch_size: int = 8
    max_beam_width: int = 1
    max_num_tokens: Optional[int] = None
    opt_num_tokens: Optional[int] = None
    max_prompt_embedding_table_size: int = 0
    kv_cache_type: KVCacheType = None
    gather_context_logits: int = False
    gather_generation_logits: int = False
    strongly_typed: bool = False
    builder_opt: Optional[int] = None
    force_num_profiles: Optional[int] = None
    profiling_verbosity: str = 'layer_names_only'
    enable_debug_output: bool = False
    max_draft_len: int = 0
    speculative_decoding_mode: SpeculativeDecodingMode = SpeculativeDecodingMode.NONE
    use_refit: bool = False
    input_timing_cache: str = None
    output_timing_cache: str = None
    lora_config: LoraConfig = field(default_factory=LoraConfig)
    auto_parallel_config: AutoParallelConfig = field(
        default_factory=AutoParallelConfig)
    weight_sparsity: bool = False
    weight_streaming: bool = False
    plugin_config: PluginConfig = field(default_factory=PluginConfig)
    use_strip_plan: bool = False
    max_encoder_input_len: int = 1  # for enc-dec DecoderModel
    use_fused_mlp: bool = False
    dry_run: bool = False
    visualize_network: bool = False

    # Since we have some overlapping between kv_cache_type, paged_kv_cache, and paged_state (later two will be deprecated in the future),
    # we need to handle it given model architecture.
    def update_kv_cache_type(self, model_architecture: str):
        paged_kv_cache_attr = 'paged_state' if model_architecture in [
            'MambaForCausalLM', 'RecurrentGemmaForCausalLM'
        ] else 'paged_kv_cache'
        assert self.plugin_config is not None
        paged_kv_cache_val = getattr(self.plugin_config, paged_kv_cache_attr)

        if self.kv_cache_type is not None:
            if paged_kv_cache_val is not None:
                assert (paged_kv_cache_val == True
                        and self.kv_cache_type == KVCacheType.PAGED) or (
                            paged_kv_cache_val == False
                            and self.kv_cache_type != KVCacheType.PAGED)
            else:
                setattr(self.plugin_config, paged_kv_cache_attr,
                        self.kv_cache_type == KVCacheType.PAGED)
        else:
            if paged_kv_cache_val is not None:
                self.kv_cache_type = KVCacheType.PAGED if paged_kv_cache_val else KVCacheType.CONTINUOUS
            else:
                self.kv_cache_type = KVCacheType.PAGED
                setattr(self.plugin_config, paged_kv_cache_attr,
                        self.kv_cache_type == KVCacheType.PAGED)

        assert self.kv_cache_type is not None and getattr(
            self.plugin_config, paged_kv_cache_attr) is not None

        def override_attri(attr_name, value):
            val = getattr(self.plugin_config, attr_name)
            if val is not None and val != value:
                logger.warning(f'Overriding {attr_name} to {value}')
            setattr(self.plugin_config, attr_name, value)

        # Init other paged kvcache attri to false. For RecurrentGemma, we only support paged_state and paged_kv_cache have
        # the same values. All other models should only consume either of the value and set other to False.
        is_recurrent_gemma = model_architecture == 'RecurrentGemmaForCausalLM'

        if paged_kv_cache_attr == 'paged_state':
            override_attri(
                'paged_kv_cache',
                getattr(self.plugin_config, paged_kv_cache_attr)
                if is_recurrent_gemma else False)
        else:
            override_attri('paged_state', False)

    @classmethod
    def from_dict(cls, config, plugin_config=None):
        max_input_len = config.pop('max_input_len')
        max_seq_len = config.pop('max_seq_len')
        max_batch_size = config.pop('max_batch_size')
        max_beam_width = config.pop('max_beam_width')
        max_num_tokens = config.pop('max_num_tokens')
        opt_num_tokens = config.pop('opt_num_tokens')
        opt_batch_size = config.pop('opt_batch_size', None)
        max_prompt_embedding_table_size = config.pop(
            'max_prompt_embedding_table_size', 0)

        kv_cache_type = KVCacheType(
            config.pop('kv_cache_type')) if 'plugin_config' in config else None
        gather_context_logits = config.pop('gather_context_logits', False)
        gather_generation_logits = config.pop('gather_generation_logits', False)
        strongly_typed = config.pop('strongly_typed', True)
        builder_opt = config.pop('builder_opt', None)
        force_num_profiles = config.pop('force_num_profiles', None)
        weight_sparsity = config.pop('weight_sparsity', False)
        profiling_verbosity = config.pop('profiling_verbosity',
                                         'layer_names_only')
        enable_debug_output = config.pop('enable_debug_output', False)
        max_draft_len = config.pop('max_draft_len', 0)
        speculative_decoding_mode = config.pop('speculative_decoding_mode',
                                               SpeculativeDecodingMode.NONE)
        use_refit = config.pop('use_refit', False)
        input_timing_cache = config.pop('input_timing_cache', None)
        output_timing_cache = config.pop('output_timing_cache', None)
        lora_config = LoraConfig.from_dict(config.get('lora_config', {}))
        auto_parallel_config = AutoParallelConfig.from_dict(
            config.get('auto_parallel_config', {}))
        max_encoder_input_len = config.pop('max_encoder_input_len', 1024)
        weight_streaming = config.pop('weight_streaming', False)

        use_strip_plan = config.pop('use_strip_plan', False)

        if plugin_config is None:
            plugin_config = PluginConfig()
        if "plugin_config" in config.keys():
            plugin_config.update_from_dict(config["plugin_config"])

        dry_run = config.pop('dry_run', False)
        visualize_network = config.pop('visualize_network', False)

        return cls(
            max_input_len=max_input_len,
            max_seq_len=max_seq_len,
            max_batch_size=max_batch_size,
            max_beam_width=max_beam_width,
            max_num_tokens=max_num_tokens,
            opt_num_tokens=opt_num_tokens,
            opt_batch_size=opt_batch_size,
            max_prompt_embedding_table_size=max_prompt_embedding_table_size,
            kv_cache_type=kv_cache_type,
            gather_context_logits=gather_context_logits,
            gather_generation_logits=gather_generation_logits,
            strongly_typed=strongly_typed,
            builder_opt=builder_opt,
            force_num_profiles=force_num_profiles,
            profiling_verbosity=profiling_verbosity,
            enable_debug_output=enable_debug_output,
            max_draft_len=max_draft_len,
            speculative_decoding_mode=speculative_decoding_mode,
            use_refit=use_refit,
            input_timing_cache=input_timing_cache,
            output_timing_cache=output_timing_cache,
            lora_config=lora_config,
            auto_parallel_config=auto_parallel_config,
            use_strip_plan=use_strip_plan,
            max_encoder_input_len=max_encoder_input_len,
            weight_sparsity=weight_sparsity,
            weight_streaming=weight_streaming,
            plugin_config=plugin_config,
            dry_run=dry_run,
            visualize_network=visualize_network)

    @classmethod
    def from_json_file(cls, config_file, plugin_config=None):
        with open(config_file) as f:
            config = json.load(f)
            return BuildConfig.from_dict(config, plugin_config=plugin_config)

    def to_dict(self):
        output = copy.deepcopy(self.__dict__)
        output['plugin_config'] = output['plugin_config'].to_dict()
        output['lora_config'] = output['lora_config'].to_dict()
        output['auto_parallel_config'] = output['auto_parallel_config'].to_dict(
        )
        return output

    def update_from_dict(self, config: dict):
        for name, value in config.items():
            if not hasattr(self, name):
                raise AttributeError(
                    f"{self.__class__} object has no attribute {name}")
            setattr(self, name, value)

    def update(self, **kwargs):
        self.update_from_dict(kwargs)


class EngineConfig:

    def __init__(self, pretrained_config: 'PretrainedConfig',
                 build_config: 'BuildConfig', version: str):
        self.pretrained_config = pretrained_config
        self.build_config = build_config
        self.version = version

    @classmethod
    def from_json_file(cls, config_file):
        with open(config_file) as f:
            config = json.load(f)
            return cls(PretrainedConfig.from_dict(config['pretrained_config']),
                       BuildConfig.from_dict(config['build_config']),
                       config['version'])

    def to_dict(self):
        build_config = self.build_config.to_dict()
        build_config.pop('dry_run', None)  # Not an Engine Characteristic
        build_config.pop('visualize_network',
                         None)  # Not an Engine Characteristic
        return {
            'version': self.version,
            'pretrained_config': self.pretrained_config.to_dict(),
            'build_config': build_config,
        }


class Engine:

    def __init__(
        self,
        config: EngineConfig,
        engine: Union[trt.IHostMemory, None],
        managed_weights: list[tuple[str, np.ndarray]] = None,
    ):
        self.config = config
        self.engine = engine
        self.managed_weights = managed_weights

    def save(self, engine_dir: str):
        os.makedirs(engine_dir, exist_ok=True)
        lora_config = self.config.build_config.lora_config
        lora_dirs = lora_config.lora_dir
        root_lora_dir = os.path.join(engine_dir, 'lora')
        if len(lora_dirs) > 0:
            os.makedirs(root_lora_dir, exist_ok=True)
            for index, lora_dir in enumerate(lora_dirs):
                if lora_config.lora_ckpt_source == 'hf':
                    target_lora_dir = f"{root_lora_dir}/{index}"
                    os.makedirs(target_lora_dir, exist_ok=True)
                    shutil.copy2(os.path.join(lora_dir, 'adapter_config.json'),
                                 target_lora_dir)
                    weight_file = os.path.join(lora_dir, 'adapter_model.bin')
                    if os.path.exists(weight_file):
                        shutil.copy2(weight_file, target_lora_dir)
                    weight_file = os.path.join(lora_dir,
                                               'adapter_model.safetensors')
                    if os.path.exists(weight_file):
                        shutil.copy2(weight_file, target_lora_dir)
                    lora_config.lora_dir[index] = f"lora/{index}"
                elif lora_config.lora_ckpt_source == 'nemo':
                    target_lora_file = f"{root_lora_dir}/{index}.nemo"
                    shutil.copyfile(lora_dir, target_lora_file)
                    lora_config.lora_dir[index] = f"lora/{index}.nemo"
        else:
            if os.path.exists(root_lora_dir) and os.path.isdir(root_lora_dir):
                shutil.rmtree(root_lora_dir)
        if self.config.pretrained_config.mapping.rank == 0:
            with open(os.path.join(engine_dir, 'config.json'),
                      "w",
                      encoding="utf-8") as f:
                json.dump(self.config.to_dict(), f, indent=4, cls=ConfigEncoder)
        if self.engine is not None:
            serialize_engine(
                self.engine,
                os.path.join(
                    engine_dir,
                    f'rank{self.config.pretrained_config.mapping.rank}.engine'))
        if self.managed_weights is not None and len(self.managed_weights) > 0:
            fn = os.path.join(
                engine_dir,
                f'rank{self.config.pretrained_config.mapping.rank}_managed_weights.safetensors'
            )
            serialize_managed_weights(self.managed_weights, fn)

    @classmethod
    def from_dir(cls, engine_dir: str, rank: int = 0):
        with open(os.path.join(engine_dir, f'rank{rank}.engine'), 'rb') as f:
            engine_buffer = f.read()

        config = EngineConfig.from_json_file(
            os.path.join(engine_dir, 'config.json'))
        config.pretrained_config.set_rank(rank)

        return cls(config, engine_buffer)


def get_engine_version(engine_dir: str) -> Union[None, str]:
    engine_dir = Path(engine_dir)
    config_path = engine_dir / "config.json"
    with open(config_path, 'r') as f:
        config = json.load(f)

    if 'version' not in config:
        return None

    return config['version']


def optimize_model_with_config(model: PretrainedModel,
                               build_config: BuildConfig):
    use_auto_parallel = build_config.auto_parallel_config.enabled
    gemm_swiglu_plugin = build_config.plugin_config.gemm_swiglu_plugin
    if gemm_swiglu_plugin:
        if not build_config.use_fused_mlp:
            raise RuntimeError(
                "GemmSwiGLU plugin requires --use_fused_mlp flag")
        if gemm_swiglu_plugin not in ["fp8"]:
            raise RuntimeError(
                f"GemmSwiGLU plugin currently has limited support: fp8 only, "
                f"got: {gemm_swiglu_plugin}")

    if build_config.plugin_config.lora_plugin is not None:
        model.use_lora(build_config.lora_config)

    is_enc_dec = model.config.architecture in ["EncoderModel", "DecoderModel"]
    model = optimize_model(
        model,
        use_ootb_moe=build_config.plugin_config.moe_plugin is None,
        use_fused_mlp=(build_config.use_fused_mlp and not is_enc_dec
                       and not use_auto_parallel),
        gemm_swiglu_plugin_dtype=gemm_swiglu_plugin,
        use_fused_rg_lru=model.config.architecture
        in ["RecurrentGemmaForCausalLM"],
        use_unfused_qkv_gemm=use_auto_parallel,
        use_prompt_tuning=(build_config.max_prompt_embedding_table_size > 0),
        use_lora=build_config.plugin_config.lora_plugin is not None,
        max_lora_rank=build_config.lora_config.max_lora_rank,
        use_fp8_context_fmha=(
            model.config.quantization.quant_algo == QuantAlgo.FP8
            and build_config.plugin_config.use_fp8_context_fmha),
    )

    if is_enc_dec:
        model.precompute_relative_attention_bias(build_config)
    return model


def _init_max_seq_len(model_config, build_config):
    """
    If max_seq_len is not specified, set it to max_position_embeddings * rotary_factor
    Additional checks to ensure max_seq_len, max_input_len, and max_num_tokens have valid values.
    """
    # Extract rotary scaling which will be used for checks and default value of max_seq_len
    rotary_scaling = getattr(model_config, "rotary_scaling", None)
    if rotary_scaling is not None:
        rotary_type = rotary_scaling.get('type',
                                         rotary_scaling.get('rope_type'))
        rotary_factor = rotary_scaling.get('factor',
                                           1.0) if rotary_type != 'su' else 1
    else:
        rotary_factor = 1

    if model_config.architecture == "EncoderModel":
        if build_config.max_seq_len is None:
            build_config.max_seq_len = build_config.max_input_len
            logger.info(
                f'max_seq_len is not specified for EncoderModel, using --max_input_len.'
            )
        assert build_config.max_input_len == build_config.max_seq_len, f"EncoderModel should have same --max_input_len ({build_config.max_input_len}) and --max_seq_len ({build_config.max_seq_len})."

    if build_config.max_seq_len is None:
        # Step 1: Find the upper bound of max_seq_len
        deduced_max_seq_len = 2048
        if model_config.max_position_embeddings is not None:
            deduced_max_seq_len = model_config.max_position_embeddings

        # Step 2: Scale max_seq_len with rotary scaling
        if rotary_factor != 1:
            deduced_max_seq_len = math.ceil(deduced_max_seq_len * rotary_factor)
            logger.warning(
                f'max_seq_len is scaled to {deduced_max_seq_len} by rotary scaling {rotary_factor}'
            )

        # Step 3: Assign the new max_seq_len
        build_config.max_seq_len = int(deduced_max_seq_len)
        logger.info(
            f'max_seq_len is not specified, using deduced value {deduced_max_seq_len}'
        )
    else:
        if not build_config.plugin_config.streamingllm and model_config.max_position_embeddings is not None \
            and model_config.position_embedding_type != PositionEmbeddingType.relative:
            if build_config.max_seq_len > model_config.max_position_embeddings * rotary_factor:
                logger.warning(
                    f'max_seq_len {build_config.max_seq_len} is larger than max_position_embeddings {model_config.max_position_embeddings} * rotary scaling {rotary_factor}, '
                    'the model accuracy might be affected')

    if build_config.max_input_len > build_config.max_seq_len:
        logger.warning(
            f'max_input_len is {build_config.max_input_len} is larger than max_seq_len {build_config.max_seq_len}, clipping it to max_seq_len'
        )
        build_config.max_input_len = build_config.max_seq_len

    # Check and may modify max_num_tokens and opt_num_tokens (need to happen after max_seq_len is deduced)
    max_num_tokens, opt_num_tokens = check_max_num_tokens(
        max_num_tokens=build_config.max_num_tokens,
        opt_num_tokens=build_config.opt_num_tokens,
        max_batch_size=build_config.max_batch_size,
        max_input_len=build_config.max_input_len,
        max_seq_len=build_config.max_seq_len,
        max_beam_width=build_config.max_beam_width,
        remove_input_padding=build_config.plugin_config.remove_input_padding,
        enable_context_fmha=build_config.plugin_config.context_fmha,
        tokens_per_block=build_config.plugin_config.tokens_per_block,
        multiple_profiles=build_config.plugin_config.multiple_profiles,
    )
    build_config.max_num_tokens, build_config.opt_num_tokens = max_num_tokens, opt_num_tokens

    if build_config.plugin_config.remove_input_padding and build_config.plugin_config.context_fmha:
        if build_config.max_input_len:
            logger.warning(
                'padding removal and fMHA are both enabled, max_input_len is not required and will be ignored'
            )
    else:
        assert build_config.max_input_len is not None, 'padding removal and fMHA aren\'t both enabled, max_input_len is required'
        if build_config.max_seq_len:
            assert build_config.max_input_len <= build_config.max_seq_len, 'max_input_len should not be larger than max_seq_len'


def serialize_managed_weights(managed_weights: list[tuple[str, np.ndarray]],
                              path: str | Path,
                              metadata=None) -> None:
    header = {}
    if metadata is not None:
        header["__metadata__"] = metadata
    begin = 0
    for name, value in managed_weights:
        size = value.size * value.itemsize
        if value.dtype == np.float32:
            dtype = "F32"
        elif value.dtype == np.float16:
            dtype = "F16"
        elif value.dtype == np_bfloat16:
            dtype = "BF16"
        elif value.dtype == np_float8:
            dtype = "F8_E4M3"
        elif value.dtype == np.int64:
            dtype = "I64"
        elif value.dtype == np.int32:
            dtype = "I32"
        else:
            raise RuntimeError(f"Unsupported dtype: {value.dtype}")
        header[name] = {
            "dtype": dtype,
            "shape": value.shape,
            "data_offsets": [begin, begin + size],
        }
        begin += size

    header_json = json.dumps(header)
    header_json_len = len(header_json)
    with open(path, "wb") as f:
        logger.info(
            f"Serializing {len(managed_weights)} managed weights to {path}...")
        f.write(header_json_len.to_bytes(8, byteorder="little"))
        f.write(header_json.encode())
        for name, value in managed_weights:
            logger.debug(f"Serializing managed weight: {name}")
            buf = value.tobytes()
            f.write(buf)


def build(model: PretrainedModel,
          build_config: BuildConfig,
          return_build_config: bool = False) -> Engine | BuildConfig:
    '''Build engine from given model and optimization options specified in the build_config
       WARNING: this function may change the given \p model object state in some optimization passes
       to avoid cloning a model since normally the LLM models consumes large memory.
       Create a new fresh model object if you need to build with different options.

    '''
    # avoid changing the input config
    build_config = copy.deepcopy(build_config)
    build_config.plugin_config.dtype = model.config.dtype
    build_config.update_kv_cache_type(model.config.architecture)

    _init_max_seq_len(model.config, build_config)

    if model.config.quantization.quant_algo == QuantAlgo.FP8 or \
            model.config.quantization.kv_cache_quant_algo == QuantAlgo.FP8:
        build_config.strongly_typed = True

    if hasattr(model.config, 'max_draft_len'):
        build_config.max_draft_len = model.config.max_draft_len
        if build_config.speculative_decoding_mode != SpeculativeDecodingMode.MEDUSA:
            logger.warning(
                'speculative_decoding_mode is not Medusa for Medusa model. Overwriting speculative_decoding_mode'
            )
        build_config.speculative_decoding_mode = SpeculativeDecodingMode.MEDUSA

    if hasattr(model.config, 'redrafter_num_beams') and hasattr(
            model.config, 'redrafter_draft_len_per_beam'):
        build_config.max_draft_len = model.config.redrafter_num_beams * model.config.redrafter_draft_len_per_beam
        if build_config.speculative_decoding_mode != SpeculativeDecodingMode.EXPLICIT_DRAFT_TOKENS:
            logger.warning(
                'speculative_decoding_mode is not EXPLICIT_DRAFT_TOKENS for ReDrafter model. Overwriting speculative_decoding_mode'
            )
        build_config.speculative_decoding_mode = SpeculativeDecodingMode.EXPLICIT_DRAFT_TOKENS

    if build_config.speculative_decoding_mode != SpeculativeDecodingMode.NONE:
        logger.info(
            f'Increasing max_seq_len ({build_config.max_seq_len}) '
            f'by max_draft_len ({build_config.max_draft_len}) '
            'to account for speculative decoding implementation specifics. '
            'Maximum number of generated tokens remains the same. '
            f'New max_seq_len is set to {build_config.max_seq_len + build_config.max_draft_len}'
        )
        build_config.max_seq_len += build_config.max_draft_len

    if build_config.speculative_decoding_mode != SpeculativeDecodingMode.NONE:
        num_tokens = build_config.max_batch_size * (build_config.max_draft_len +
                                                    1)
        if build_config.max_num_tokens < num_tokens:
            logger.info(
                f'max_num_tokens ({build_config.max_num_tokens}) is smaller than '
                'max_batch_size * (max_draft_len + 1) = '
                f'({build_config.max_batch_size} * ({build_config.max_draft_len} + 1)). '
                f'New max_num_tokens is set to {num_tokens}.')
            build_config.max_num_tokens = num_tokens

    if build_config.plugin_config.use_paged_context_fmha:
        if (model.config.quant_mode.has_fp8_kv_cache()
                and not model.config.quant_mode.has_fp8_qdq()):
            raise RuntimeError(
                "FP8 Paged Context FMHA only works with fp8 quantization workflow currently."
            )
        if (model.config.quant_mode.has_fp8_kv_cache()
                and not build_config.plugin_config.use_fp8_context_fmha):
            build_config.plugin_config.use_fp8_context_fmha = True
            logger.warning(
                "FP8 Context FMHA is enabled by default to support FP8 Paged Context FMHA."
            )
        if model.config.quant_mode.has_int8_kv_cache():
            raise RuntimeError(
                "Paged Context FMHA doesn't work with int8 kv cache currently.")

    if build_config.plugin_config.manage_weights:
        if model.config.quant_mode & QuantMode.INT4_WEIGHTS or model.config.quant_mode & QuantMode.INT8_WEIGHTS:
            raise RuntimeError(
                "Managed weights is not supported with int4 or int8 weights.")

    model = optimize_model_with_config(model, build_config)

    builder = Builder()
    builder_config = builder.create_builder_config(
        precision=model.config.dtype,
        use_refit=build_config.use_refit,
        timing_cache=build_config.input_timing_cache,
        int8=(model.config.quant_mode.has_act_or_weight_quant()
              and not model.config.quant_mode.has_per_group_scaling())
        or model.config.quant_mode.has_int8_kv_cache(),
        strongly_typed=build_config.strongly_typed,
        opt_level=build_config.builder_opt,
        force_num_profiles=build_config.force_num_profiles,
        profiling_verbosity=build_config.profiling_verbosity,
        quant_mode=model.config.quant_mode,
        use_strip_plan=build_config.use_strip_plan,
        weight_sparsity=build_config.weight_sparsity,
        weight_streaming=build_config.weight_streaming,
    )

    network = builder.create_network()
    network.plugin_config = build_config.plugin_config

    use_auto_parallel = build_config.auto_parallel_config.enabled
    use_weight_only = model.config.quant_mode.is_weight_only()
    per_group = model.config.quant_mode.has_per_group_scaling()
    use_smooth_quant = model.config.quant_mode.has_act_and_weight_quant()
    use_fp8_rowwise = model.config.quant_mode.has_fp8_rowwise()
    disable_weight_only_quant_plugin = model.config.disable_weight_only_quant_plugin if hasattr(
        model.config, 'disable_weight_only_quant_plugin') else False

    if use_weight_only and not disable_weight_only_quant_plugin:
        if per_group:
            network.plugin_config.weight_only_groupwise_quant_matmul_plugin = model.config.dtype
        else:
            network.plugin_config.weight_only_quant_matmul_plugin = model.config.dtype
    if use_smooth_quant and model.config.quantization.use_plugin_sq:
        network.plugin_config.set_smooth_quant_plugins(model.config.dtype)
    if use_fp8_rowwise:
        network.plugin_config.set_fp8_rowwise_quant_plugins(model.config.dtype)
    nccl_plugin = model.config.dtype if model.config.mapping.world_size > 1 else None
    network.plugin_config.set_nccl_plugin(nccl_plugin)

    # NOTE: Please never change the build_config object after this point!
    if return_build_config:
        # Get an modified build_config that is the same as the one in the final engine dir
        return build_config

    with net_guard(network):
        # Prepare
        network.set_named_parameters(model.named_parameters())

        # Forward
        prepare_input_args = {
            "max_batch_size":
            build_config.max_batch_size,
            "max_input_len":
            build_config.max_input_len,
            "max_seq_len":
            build_config.max_seq_len,
            "use_cache":
            build_config.kv_cache_type != KVCacheType.DISABLED,
            "max_beam_width":
            build_config.max_beam_width,
            "max_num_tokens":
            build_config.max_num_tokens,
            "opt_num_tokens":
            build_config.opt_num_tokens,
            "prompt_embedding_table_size":
            build_config.max_prompt_embedding_table_size,
            "max_draft_len":
            build_config.max_draft_len,
            "speculative_decoding_draft_tokens_external":
            build_config.speculative_decoding_mode ==
            SpeculativeDecodingMode.DRAFT_TOKENS_EXTERNAL,
            "gather_context_logits":
            build_config.gather_context_logits,
            "gather_generation_logits":
            build_config.gather_generation_logits,
            "lora_target_modules":
            build_config.lora_config.lora_target_modules
        }

        if model.config.architecture == "DecoderModel":
            prepare_input_args["max_seq_len"] = build_config.max_seq_len
            prepare_input_args[
                "max_decoder_input_len"] = build_config.max_input_len
            prepare_input_args[
                "max_encoder_input_len"] = build_config.max_encoder_input_len

        if model.config.architecture == "WhisperEncoder":

            prepare_input_args = {
                "max_batch_size": build_config.max_batch_size,
            }

        inputs = model.prepare_inputs(**prepare_input_args)
        model(**inputs)

        if build_config.enable_debug_output:
            for k, v in model.named_network_outputs():
                network._mark_output(v, k, str_dtype_to_trt(model.config.dtype))

    if model.config.architecture != "DecoderModel":
        optimize(network)

    if use_auto_parallel:
        config = build_config.auto_parallel_config
        config.builder_flags = builder_config.trt_builder_config.flags
        sharded_networks = auto_parallel(network, config)
        network = sharded_networks[model.config.mapping.rank]
        if not build_config.auto_parallel_config.debug_mode:
            mapping = network.auto_parallel_config["mapping"]
            model.config.mapping = mapping

    if build_config.visualize_network:
        with net_guard(network):
            network.to_dot(f'rank{model.config.mapping.rank}.dot')

    # Network -> Engine
    managed_weights = [] if network.plugin_config.manage_weights else None
    engine = None if build_config.dry_run else builder.build_engine(
        network, builder_config, managed_weights)
    engine_config = EngineConfig(model.config, build_config, __version__)

    if build_config.output_timing_cache is not None and model.config.mapping.rank == 0:
        ok = builder.save_timing_cache(builder_config,
                                       build_config.output_timing_cache)
        assert ok, "Failed to save timing cache."

    import psutil

    # Get the current process
    current_process = psutil.Process()
    # Get resource usage for the current process (self)
    rusage_s = current_process.memory_info()
    # Get resource usage for all child processes
    children = current_process.children(recursive=True)
    rusage_c = [child.memory_info() for child in children]
    logger.info(
        f"Build phase peak memory: {rusage_s.rss / 1024 / 1024:.2f} MB, children: {sum([ru.rss for ru in rusage_c]) / 1024 / 1024:.2f} MB"
    )

    return Engine(engine_config, engine, managed_weights)