Function latency_command

tensorrt_llm/bench/benchmark/low_latency.py:190–373

Run a latency test on a TRT-LLM engine.

(
    bench_env: BenchmarkEnvironment,
    **params,
)

)
@click.pass_obj
def latency_command(
    bench_env: BenchmarkEnvironment,
    **params,
) -> None:
    """Run a latency test on a TRT-LLM engine."""
    logger.info("Preparing to run latency benchmark...")

    # Parameters from CLI
    # Model, experiment, and engine params
    options = get_general_cli_options(params, bench_env)

    # Speculative Decode Options
    medusa_choices = params.get("medusa_choices")
    # Initialize the HF tokenizer for the specified model.
    tokenizer = initialize_tokenizer(options.checkpoint_path)

    # Dataset Loading and Preparation
    with open(options.dataset_path, "r") as dataset:
        metadata, requests = create_dataset_from_stream(
            tokenizer,
            dataset,
            num_requests=options.num_requests,
            model_dir=options.checkpoint_path,
            model_type=options.model_type,
            modality=options.modality,
            max_input_seq_len_for_multimodal=options.max_input_len)

        metadata.dataset_path = options.dataset_path

    if options.modality is None:
        # Log dataset info
        # NOTE: This table is only accurate for non-multimodal models.
        #       The accurate table for multimodal models will be logged after the benchmark is done.
        logger.info(metadata.get_summary_for_print())

    # Engine configuration parsing for PyTorch backend
    kwargs = {}
    if options.backend and options.backend.lower(
    ) in ALL_SUPPORTED_BACKENDS and options.backend.lower() != "tensorrt":
        if bench_env.checkpoint_path is None:
            snapshot_download(options.model, revision=bench_env.revision)

        exec_settings = get_settings(params, metadata, bench_env.model,
                                     bench_env.checkpoint_path)
        kwargs_max_sql = options.max_seq_len or metadata.max_sequence_length
        logger.info(f"Setting PyTorch max sequence length to {kwargs_max_sql}")
        kwargs["max_seq_len"] = kwargs_max_sql
    elif options.backend.lower() == "tensorrt":
        assert options.max_seq_len is None, (
            "max_seq_len is not a runtime parameter for C++ backend")
        exec_settings, build_cfg = get_settings_from_engine(options.engine_dir)
        engine_max_seq_len = build_cfg["max_seq_len"]

        if metadata.max_sequence_length > engine_max_seq_len:
            raise RuntimeError(
                f"Engine supports a max sequence of {engine_max_seq_len}. Provided "
                "dataset contains a maximum sequence of "
                f"{metadata.max_sequence_length}. Please rebuild a new engine to"