nemo_rl.models.generation.trtllm.trtllm_http_server

OpenAI-compatible HTTP server wrapping tensorrt_llm.LLM, serving /v1/chat/completions.

Returns NeMoGym fields (prompt_token_ids, generation_token_ids, generation_log_probs). Supports Qwen3 tool calling, DeepSeekR1Parser reasoning, and prefix token splicing.

Module Contents

Functions

_build_reasoning_parser
_build_sampling_params	Build the TRT-LLM sampling params for one HTTP rollout request.
create_app	Build a FastAPI application backed by llm (tensorrt_llm.LLM).
_resolve_tool_parser_name	Resolve the configured parser or infer it from the model.
_build_tool_parser	Instantiate a TRT-LLM tool parser by registered name.
_make_parse_tool_calls	Return a tool-call parser bound to a specific parser instance.
_to_int_ids	Coerce chat-template output to a flat list[int].
_build_prompt_token_ids	Convert chat messages to token IDs via apply_chat_template (full retokenisation each turn).
_compute_splice_inputs	Return preserved and rendered token IDs for the on-policy prefix splice.
start_server	Start the HTTP server in a daemon thread and return (thread, base_url, server).

Data

logger

API

nemo_rl.models.generation.trtllm.trtllm_http_server.logger

‘getLogger(…)’

nemo_rl.models.generation.trtllm.trtllm_http_server._build_reasoning_parser(

name: str,

chat_template_kwargs: dict[str, Any],

) → Any

nemo_rl.models.generation.trtllm.trtllm_http_server._build_sampling_params(

sampling_params_cls: Any,

*,

sampling_config: dict[str, Any],

stop_token_ids: list[int] | None,

max_tokens: int,

) → Any

Build the TRT-LLM sampling params for one HTTP rollout request.

Mirrors the direct generate() path (TrtllmAsyncGenerationWorkerImpl._build_sampling_params) so both paths sample from the same distribution for a given generation config.

Parameters:

• sampling_params_cls – tensorrt_llm.SamplingParams, injected so this helper stays importable and testable without the TRT-LLM runtime.

• sampling_config – NeMo-RL generation config (temperature / top_p / top_k).

• stop_token_ids – Extra stop tokens from the generation config, if any.

• max_tokens – Output cap for this request, already clamped to the context window.

Returns:

A SamplingParams instance to hand to llm.generate_async.

nemo_rl.models.generation.trtllm.trtllm_http_server.create_app(

llm: Any,

tokenizer: Any,

model_name: str,

max_seq_len: int,

sampling_config: dict[str, Any],

stop_token_ids: list[int] | None = None,

default_chat_template_kwargs: dict[str, Any] | None = None,

tool_parser: str | None = None,

reasoning_parser: str | None = None,

) → fastapi.FastAPI

Build a FastAPI application backed by llm (tensorrt_llm.LLM).

nemo_rl.models.generation.trtllm.trtllm_http_server._resolve_tool_parser_name(

configured_name: str | None,

model_name: str,

) → str

Resolve the configured parser or infer it from the model.

nemo_rl.models.generation.trtllm.trtllm_http_server._build_tool_parser(name: str) → Any

Instantiate a TRT-LLM tool parser by registered name.

nemo_rl.models.generation.trtllm.trtllm_http_server._make_parse_tool_calls(tool_parser_instance: Any) → Any

Return a tool-call parser bound to a specific parser instance.

nemo_rl.models.generation.trtllm.trtllm_http_server._to_int_ids(enc: Any) → list[int]

Coerce chat-template output to a flat list[int].

nemo_rl.models.generation.trtllm.trtllm_http_server._build_prompt_token_ids(

messages: list[dict[str, Any]],

tokenizer: Any,

*,

tools: list[dict[str, Any]] | None = None,

default_template_kwargs: dict[str, Any] | None = None,

) → list[int]

Convert chat messages to token IDs via apply_chat_template (full retokenisation each turn).

Full retokenisation avoids the gen_token_ids double-counting bug (~4000 tok/turn explosion that exhausted context at turn ~32 when prefix accumulation was used).

nemo_rl.models.generation.trtllm.trtllm_http_server._compute_splice_inputs(

raw_messages: list[dict[str, Any]],

conversation: list[dict[str, Any]],

tokenizer: Any,

tools: list[dict[str, Any]] | None,

default_template_kwargs: dict[str, Any],

) → tuple[list[int], list[int]]

Return preserved and rendered token IDs for the on-policy prefix splice.

nemo_rl.models.generation.trtllm.trtllm_http_server.start_server(

llm: Any,

tokenizer: Any,

model_name: str,

max_seq_len: int,

sampling_config: dict[str, Any],

stop_token_ids: list[int] | None = None,

host: str = '0.0.0.0',

port: int = 0,

default_chat_template_kwargs: dict[str, Any] | None = None,

tool_parser: str | None = None,

reasoning_parser: str | None = None,

) → tuple[threading.Thread, str, Any]

Start the HTTP server in a daemon thread and return (thread, base_url, server).