fix(trtllm): Align HTTP server sampling params, use generate_async - #3537
Merged
Conversation
shuyixiong
added
the
CI:Lfast
shuyixiong
force-pushed
the
shuyix/trtllm_server_fix
branch
from
August 7, 2026 13:38
0aba997 to
612f958
Compare
Contributor
Author
|
/ok to test 612f958 |
The HTTP server left logprobs in the default per-token dict format while the direct generate() path (_build_sampling_params) already asked for the flat one. TRT-LLM normalizes logprobs=True to 0 before validating logprobs_simple_format, so the combination is accepted, and the response handler already accepts both shapes -- it just stops allocating a dict per generated token. The same divergence hid a second one: the HTTP server built its SamplingParams without top_k, so a configured generation.top_k was silently ignored for every NeMo-Gym rollout while the direct path applied it, leaving the two paths sampling from different distributions on the same config. Pass it through, mapping an unset value to 0 the way the direct path does, since that is how TRT-LLM spells "no top-k restriction". top_k now also participates in the request/config equality check that already guarded temperature and top_p, so a request cannot ask for a sampling profile the generation config did not specify. The lookups there move to .get() so a server whose sampling_config predates this key does not raise KeyError. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
The handler ran the blocking llm.generate() on a worker thread via asyncio.to_thread, so every concurrent rollout request consumed a thread from the default executor while it waited on the engine. Multi-turn SWE rollouts keep hundreds of requests in flight, which is far more than that pool is sized for, so requests queued on threads rather than on the engine's own scheduler. generate_async is the engine's native awaitable and needs no thread: the request goes straight onto the executor's queue and the coroutine parks on its future. It takes a single prompt instead of a batch, so the one-element list and the outputs[0] unwrap go away, and asyncio is no longer used anywhere in this module. RequestError still surfaces the same way, leaving the context-length 400 path unchanged. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
shuyixiong
force-pushed
the
shuyix/trtllm_server_fix
branch
from
August 7, 2026 14:42
612f958 to
bf2602e
Compare
Contributor
Author
|
/ok to test bf2602e |
yuki-97
reviewed
Aug 10, 2026
yuki-97
left a comment
yuki-97
left a comment
Contributor
There was a problem hiding this comment.
@shuyixiong thanks for the fix! left some minor comments.
Contributor
Author
|
/ok to test 889aaed |
Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
LLM._prepare_sampling_params already calls sampling_params._setup(self.tokenizer, ...) which sets end_id = tokenizer.eos_token_id when None, making the explicit AutoConfig lookup redundant. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
The HTTP server built its SamplingParams without stop_token_ids while the direct generate() path passed them, so a configured generation.stop_token_ids only reached the engine on one of the two paths. TRT-LLM's SamplingParams._setup appends generation_config.eos_token_id to stop_token_ids rather than replacing it, so passing them is additive and leaves the model's own EOS handling intact. Extract the construction into a module-level _build_sampling_params that takes the SamplingParams class as an argument, mirroring the direct path's method of the same name. It was previously inline inside create_app's request closure, which made the sampling params untestable without standing up the whole app. The sampling_config lookups also drop .get() for subscripts: top_k is a required GenerationConfig key and the sole construction site fills all three, so .get() would turn a missing key into a silent "no top-k restriction" instead of failing. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: shuyixiong <219646547+shuyixiong@users.noreply.github.com>
shuyixiong
force-pushed
the
shuyix/trtllm_server_fix
branch
from
August 11, 2026 01:56
889aaed to
4c56846
Compare
yuki-97
approved these changes
Aug 11, 2026
Contributor
|
/ok to test 4c56846 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Make the TRT-LLM HTTP rollout path sample the same way the direct
generate()path already does, and serve it without a worker thread.TrtllmAsyncGenerationWorkerreaches the engine two ways:generate()/generate_async()call_build_sampling_params()directly, while NeMo-Gym rollouts go through the HTTP server, which built its ownSamplingParams. The two had drifted:top_klogprobs_simple_format=Trueasyncio.to_thread(llm.generate, ...)Before your PR is "Ready for review"
Pre checks:
Additional Information