Method from_dir

tensorrt_llm/runtime/model_runner.py:639–782

Create a ModelRunner instance from an engine directory. Args: engine_dir (str): The directory that contains the serialized engine files and config files. max_output_len (Optional[int]): max_output_len, this arg might be availa

(
        cls,
        engine_dir: str,
        *,
        max_output_len: Optional[int] = None,
        lora_dir: Optional[List[str]] = None,
        rank: int = 0,
        debug_mode: bool = False,
        lora_ckpt_source: str = "hf",
        medusa_choices: List[List[int]] = None,
        stream: torch.cuda.Stream = None,
        gpu_weights_percent: float = 1,
        enable_context_fmha_fp32_acc: Optional[bool] = None,
        multi_block_mode: Optional[bool] = None,
        fail_fast_on_attention_window_too_large: bool = False,
    )

    @classmethod
    def from_dir(
        cls,
        engine_dir: str,
        *,
        max_output_len: Optional[int] = None,
        lora_dir: Optional[List[str]] = None,
        rank: int = 0,
        debug_mode: bool = False,
        lora_ckpt_source: str = "hf",
        medusa_choices: List[List[int]] = None,
        stream: torch.cuda.Stream = None,
        gpu_weights_percent: float = 1,
        enable_context_fmha_fp32_acc: Optional[bool] = None,
        multi_block_mode: Optional[bool] = None,
        fail_fast_on_attention_window_too_large: bool = False,
    ) -> 'ModelRunner':
        """
        Create a ModelRunner instance from an engine directory.

        Args:
            engine_dir (str):
                The directory that contains the serialized engine files and config files.
            max_output_len (Optional[int]):
                max_output_len, this arg might be available only when loading time, generate will still to check when disable_kv_cache is enabled.
            lora_dir (Optional[List[str]]):
                The directories that contain LoRA weights.
            rank (int):
                The runtime rank id.
            debug_mode (bool):
                Whether or not to turn on the debug mode.
            medusa_choices (List[List[int]]):
                Medusa choices to use when in Medusa decoding
            stream (torch.cuda.Stream):
                Stream to use.
            multi_block_mode (bool):
                Whether to distribute the work across multiple CUDA thread-blocks on the GPU for masked MHA kernel.
            fail_fast_on_attention_window_too_large (bool):
                Exit with runtime error when attention window is too large to fit even a single sequence in the KV cache.
                Note: This parameter is only applicable to C++ runtime (ModelRunnerCpp).
        Returns:
            ModelRunner: An instance of ModelRunner.
        """
        engine_version = get_engine_version(engine_dir)
        profiler.start('load tensorrt_llm engine')
        # the old engine format
        if engine_version is None:
            engine_dir = Path(engine_dir)
            config_path = engine_dir / "config.json"
            model_config, other_config = read_config(config_path)
            world_size = other_config.pop('world_size')
            tp_size = other_config.pop('tp_size')
            pp_size = other_config.pop('pp_size')
            max_batch_size = other_config.pop('max_batch_size')
            max_input_len = other_config.pop('max_input_len')
            max_output_len = other_config.pop('max_output_len')
            max_beam_width = other_config.pop('max_beam_width')
            runtime_mapping = Mapping(world_size=world_size,
                                      rank=rank,