feat(trtllm): support agentic training for SWE rollout

[hchings]

What does this PR do ?

Add support of agentic training for SWE with TRTLLM rollout.
This MR builds on top of #2420 (merged).

[‼️ Notes for reviewers ‼️]

• This MR only ensure convergence of trtllm aligns w/ vllm's. Perf optimizations will be in subsequent MRs.

Scope of testing:

1. Qwen3-30B-Thinking, 16 nodes X H100s

• TRTLLM wandb, using the recipe provided by Rayen. Please see pinned run w/ suffix eos-fix-0630-r3.
• Convergence aligned with BIn Hu's run with vllm. vllm wandb, verified by Jia in early July.

2. Qwen3.5-35B-A3B 4 nodes X GB200/GB300

• TRTLLM wandb, 07/22
• vllm wandb
• Convergence related curves aligned between both backends in the first 20+ steps.

All above will be rerun to verify convergence again once MR is in its final stage.

TODO

• Formalize env packaging for TRTLLM
• Add CI test(s)
• vLLM<->TRT-LLM chat-completions parity test
• Retest convergence as sanity after addressing comments

Issues

TRTLLM-13972

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuki-97]

thanks @hchings , left some comments.

this PR support gym path, it's worth to add a functional test and nightly test for this, you can take 3eedbc6 as an example for nightly test.

Reviewed via hchings#4 (delta against shuyix/trtllm), scoped to the agentic-SWE + TRT-LLM HTTP server changes. Scripts and recipe YAMLs skipped per the PR note. Comments target the shared code paths.

[terrykong]

Thanks @hchings — the TRT-LLM agentic path is a valuable addition, and the convergence-parity scoping (with perf deferred) is a reasonable split. This pass was reviewed via hchings#4 (the true delta against shuyix/trtllm, since #3130 is stacked on the not-yet-merged #2420), scoped to the 14 files that PR adds. Same head SHA, so comments anchor here.

Rebase needed: the PR currently has merge conflicts against main (mergeable: CONFLICTING) — please rebase on main (or the refreshed #2420 base) and resolve before merge.

A few non-blocking notes (details in the inline comments where actionable):

• Verified positive: the new test_replace_prefix_tokens_qwen3_think_shift_picks_assistant_eos_not_user_eos is a genuine regression test — we executed both the old positional and new EOS-count algorithms on its inputs; the old one drops the intervening user turn, the new one preserves it.
• We also verified the open logprob question on the earlier review thread is not a bug on this code path (evidence permalinks in the thread reply) — no change needed there.
• The pure helpers in trtllm_http_server.py (_parse_tool_calls, _to_int_ids, _compute_splice_inputs, ...) currently have no unit coverage; they are CPU-testable with a stub tokenizer. We have a verified-passing 8-test suite for them — happy to share it on this PR if useful.
• Nightly/functional coverage for the gym+trtllm path is still on your TODO — +1 to landing at least one entry with (or right after) this PR.
• For the perf follow-up PRs: asyncio.to_thread(llm.generate, ...) uses asyncio's default executor (max ~32 threads), which caps concurrent HTTP generations well below max_batch_size: 256 — worth revisiting when you optimize throughput (e.g. a dedicated executor or llm.generate_async).

Generated by Claude Code

[terrykong]

trtllm_http_server.py:41

Suggestion: add a vLLM<->TRT-LLM chat-completions parity test. This server independently re-implements prompt construction (_build_prompt_token_ids/_compute_splice_inputs), tool parsing (_parse_tool_calls), reasoning parsing, and EOS-stripping — the vLLM server (vllm_worker_async.py) does all of these differently, and the two share only replace_prefix_tokens. Nothing currently guarantees they emit identical tokens for the same input: the existing test_vllm_http_server_correct_merged_tokens_matches_baseline checks vLLM-server vs vLLM-engine only, and tests/unit/models/generation/trtllm/ has no HTTP-server test at all.

Action: add a greedy (temperature=0.0) parity test — same small model (e.g. Qwen3-0.6B, TP=1) and the same multi-turn scenario through each backend's /v1/chat/completions, asserting bitwise-identical prompt_token_ids, generation_token_ids, and content (the exact fields NeMo-Gym reads back — nemo_gym.py#L353). The scenario should include a tool-call turn (request with tools=, assistant emits a <tool_call>, then a tool result message and a follow-up turn) — that exercises the highest-drift surfaces at once: tool parsing/serialization, finish_reason="tool_calls", and the prefix splice across the tool boundary. Since vllm/trtllm live in separate extras and CI lanes, the practical shape is a committed golden of (prompt_token_ids, generation_token_ids, content, tool_calls) that each backend's lane asserts against (TRT-LLM side as tests/unit/models/generation/trtllm/test_trtllm_http_server.py with pytestmark = pytest.mark.trtllm), or alternatively a functional/L1 nightly launching both — which would also cover the "Add CI test(s)" TODO. SGLang is out of scope here: it's driven token-in/token-out and never hits this server layer (openai_server_utils.py#L20).

[terrykong]

@bxyu-nvidia could i get another set of eyes on the subtle change logic change here?

[cmunley1]

did you measure logp error with stripped thinking blocks (Qwen3)? wouldn't it be high without something like an attn mask in this case if you are recomputing?

[bxyu-nvidia]

can we clarify that we count by EOS to enable training Qwen 3 reasoning models without needing to modify their chat template to prevent removing the interleaved thinking?

[joyang-nv]

/ok to test c1919ff

[joyang-nv]

/ok to test c298a42

[joyang-nv]

/ok to test 8fbb6ff