Method _get_sampling_config

tensorrt_llm/sampling_params.py:454–484

(self)

        return list(zip(stop_reasons, stop_words))

    def _get_sampling_config(self) -> tllme.SamplingConfig:
        # A map from the SamplingConfig fields of the LLM API to their
        # corresponding field names of the Executor of TRT-LLM C++ runtime.
        # In sampling, there is no parameter that directly matches 'best_of',
        # so outputs must be trimmed during postprocessing.
        #               |     LLM API     |    TRT-LLM Executor    |
        # --------------|-----------------|------------------------|
        # | Beam search | use_beam_search | beam_width > 1         |
        # | Beam search | n               | num_return_sequences   |
        # | Beam search | best_of         | beam_width             |
        # |-------------|-----------------|------------------------|
        # | Sampling    | use_beam_search | beam_width == 1        |
        # | Sampling    | n               | num_return_sequences   |
        # | Sampling    | best_of         | no corresponding param |
        fields = {f for f in dir(tllme.SamplingConfig) if not f.startswith("__")}
        unmatched_params = [
            "num_return_sequences",
            "beam_width",
            "n",
            "best_of",
            "use_beam_search",
        ]
        llmapi_to_rt_param_map = {f: getattr(self, f) for f in fields if f not in unmatched_params}
        if self.use_beam_search:
            llmapi_to_rt_param_map["num_return_sequences"] = self.n
            llmapi_to_rt_param_map["beam_width"] = self.best_of
        else:
            llmapi_to_rt_param_map["num_return_sequences"] = self.best_of
            llmapi_to_rt_param_map["beam_width"] = 1

        return tllme.SamplingConfig(**llmapi_to_rt_param_map)

    def _get_output_config(self, is_pytorch_backend: bool = False) -> tllme.OutputConfig:
        sampling_param_fields = set(dir(SamplingParams))