Merge #1557 chat transport into #1787: dual-transport SGLang server

[linnanwang]

Hi @Kh4L — this rebuilds NVIDIA-NeMo#1557 on top of your branch rather than merging with it. Deliberately small: 4 commits, 9 files, +512/−80, fast-forwardable onto sglang-splice-fix (no merge commit).

Formally, linnan_sglang_adaptor is not an ancestor here. Your /generate machinery is the base, and NVIDIA-NeMo#1557's /generate implementation is discarded wholesale — after comparing the two properly, yours won nearly every overlap.

The framing is your own README:

"Gym's public Chat-Completions response contract does not guarantee exact sampled integer token IDs... If that endpoint gains a stable token-ID/logprob contract, the transport can change without changing the session-splice or context-overflow rules below."

That contract shipped. Chat-side TITO (sgl-project/sglang#23751, merged Jun 4) is in the stock 0.5.13 tree — v0.5.13 was cut Jun 13; compare v0.5.13...e03dfa8 gives ahead_by=0. This is the transport change you anticipated, and exactly as you predicted, your splice and overflow rules carry over untouched.

ProRL's patch_sglang_0513_token_metadata.sh is not required — its own header says it exposes the metadata "when logprobs=true", i.e. it only makes logprobs imply flags that already exist. Setting them explicitly needs no patch and no fork.

---

What this PR changes

file	+/−	what
sglang_model/app.py	+149/−9	transport knob; chat-path preprocess + token extraction; dispatch; abort; markers from config
sglang_model/tests/test_app.py	+120	make_model defaults to transport="generate"; +8 chat tests
sglang_model/README.md	+55/−14	documents both transports
vllm_model/app.py	+65/−53	extract-method (see below)
nemo_gym/responses_converter.py	+15/−2	bug fix, independent of both PRs
tests/unit_tests/test_responses_converter.py	+51	2 regression tests
sglang_model/_logic.py	+12	unsupported_sampling_params only — your extractor untouched
configs/...for_training.yaml	+13/−2	becomes the chat config
configs/...for_training_generate.yaml	+32	new

Dispatch is the only change to your entry point (app.py:140):

if self.config.transport == "chat":
    return await super().chat_completions(request, body)
return await self._sglang_chat_completion(request, body.model_dump(exclude_unset=True))

Untouched from your branch: _full_sglang_tokenize, _sglang_followup_fragment_ids, _sglang_msg_sig, _sglang_messages_match, _sglang_rendering_sig, _build_sglang_prompt_ids, _update_sglang_session_seq, _parse_sglang_generation, _sglang_length_finish, _get_sglang_tokenizer, _get_sglang_chat_template, both tool-call regexes, and the rest of _sglang_chat_completion.

openai_utils.py, tool_parsers.py, tests/test_tool_parsers.py, tests/test_logic.py and pyproject.toml are the same blob as on your branch.

vllm_model/app.py — pure extract-method

The return_token_id_information block moves out of chat_completions into _attach_token_id_information(choice_dict, body_dict, client); the call site is the same guard plus one await. After dedent, old and new are 56 lines, textually identical (51 non-blank), and all 13 top-level AST statements match. The only live references in the block are self, choice_dict, body_dict, client — all parameters.

(My commit message says "26 executable lines byte-identical". That was under my own normalization — strip blanks/comments/docstrings — and isn't reproducible without stating the rule. The claim above is the checkable one.)

New capabilities

transport: chat (new default, sglang ≥ 0.5.13) — inherits the whole VLLMModel flow, overriding only token extraction. Templating, tool-call parsing, sampling params, auth and context-overflow all move server-side. No local tokenizer that can drift; no client-side parser to maintain; no /tokenize round-trip.

transport: generate — your path, unchanged in substance. Still required for builds predating chat-side TITO (e.g. forks serving diffusion LLMs), where the local render, the splice and the client-side parsers all earn their keep.

Overlapping functionality reused

I compared implementations, not names. Yours won 8 of 11.

concern	NVIDIA-NeMo#1557	NVIDIA-NeMo#1787 (yours)	kept	why
token/logprob extraction	branches on if otl:	+ _extract_output_ids, _validate_selected_ids	yours, now used by both transports	cross-checks recovered ids against meta_info.output_ids; typed errors on malformed containers; a legitimate zero-token generation returns ([], []) instead of raising
context overflow	cap_to_context — head-truncate then bail	_sglang_length_finish	yours	never truncates, and retains prompt_token_ids on the terminal finish_reason="length", so the turn stays attributable. Also recovers from a server-side 400 by matching the error text — NVIDIA-NeMo#1557 had no equivalent
/generate HTTP call	ng_request + hand-rolled auth header + hash(sid) % len(urls)	create_generate + _resolve_client	yours	raises through the pooled client's _raise_for_status, producing exactly the response_content-carrying error your overflow branch matches on; strips /v1 internally. My hash() sharding was also unstable across processes under PYTHONHASHSEED
prompt-id normalization	normalize_token_ids	inline in _full_sglang_tokenize	yours	your render passes chat_template=, tools=, normalize_tool_call_arguments(messages), and handles tensor returns via .tolist()
sampling params	build_sampling_params, default_max_new_tokens=1024	inline, derived from remaining window	yours, tuple widened	a flat 1024 relied on cap_to_context (deleted) to clamp; deriving from the window cannot overflow on its own
tool-call parsing	none	tool_parsers.py + _parse_sglang_generation	yours, byte-identical	nothing to merge
dependency declaration	requirements.txt (transformers>=4.44,<6)	pyproject.toml (==5.8.1)	yours	cli/setup_command.py raises if both exist; and the exact pin matters on the local-render path, where a tokenizer delta is a silent contiguity failure that a range permits
session splicing	re-renders history each turn	_build_sglang_prompt_ids + prefix validation	yours, unchanged	see below
unsupported-param reporting	unsupported_sampling_params	none	mine, reimplemented in your _logic.py	your loop silently dropped anything outside (temperature, top_p, top_k, stop)
_attach_token_id_information hook	extracted	inline in VLLMModel	mine	the vLLM step parses token_id:NNN and issues a create_tokenize round-trip; SGLang does neither. Without the seam, SGLangModel would override ~120 lines to change ~30
transport knob + config layout	two YAMLs	none	my layout, your values	uses_reasoning_parser: true and context_length: ??? are yours; NVIDIA-NeMo#1557 shipped false and a silent 4096

Verified at 0 hits in this tree: normalize_token_ids, build_sampling_params, _PASSTHROUGH_SAMPLING_PARAMS, cap_to_context, would_truncate, default_max_new_tokens, add_generation_prompt, _sglang_urls, ng_request, logprob_start_len, _chat_completions_via_generate, requirements.txt, NVIDIA-NeMo#1557's _logic.extract_*, and its entire test_logic.py.

Your splice changed my mind. On NVIDIA-NeMo#1557 I argued a multi-turn splice wasn't safe yet, because nothing validated the carried prefix against the caller's current messages — Gym agents legitimately rewrite history (browsecomp_agent runs context resets by default). Your _build_sglang_prompt_ids does validate it: _sglang_messages_match compares role/content/tool_calls signatures, _sglang_rendering_sig catches tools/template drift, and the None-returning fragment render falls back to a full render. Building on session state rather than carried message ids sidesteps the failure mode entirely. I owe NVIDIA-NeMo#1557 a correction.

Fixes applied to the generate path

• finish_reason="abort" raises on both transports. Your ladder special-cased only "length", so an abort — a server-cancelled partial — mapped to "stop" and entered the batch looking like a completed turn.
• ChatML markers moved to config (sglang_eos_markers, sglang_turn_suffix), replacing the _SGLANG_EOS_MARKERS ClassVar and the literal "<|im_end|>\n" in _sglang_eos_nl. On a Nemotron-style template (content + '\n' + '<extra_id_1>') the splice would otherwise write a malformed boundary into every follow-up prompt. Defaults unchanged, so ChatML models behave identically.
• context_length required only for generate, via a config validator — your mandatory-??? discipline is preserved for the transport that needs it.
• Wider sampling passthrough (frequency_penalty, repetition_penalty, min_p) plus a warning for what /generate can't honor.

What is NOT verified

Being explicit, since this is the training-data path:

• test_app.py has never been executed. No fastapi in my environment (ModuleNotFoundError on collection), so its 16 tests — and the 72 in vllm_model/tests/test_app.py that would empirically confirm the extract-method — are compile-checked only. Only the 25 pure tests (test_logic.py, test_tool_parsers.py) actually ran.
• Your 8 generate-path tests are what I'd most want your eyes on. make_model now defaults to transport="generate" so they still exercise your path, but I could not run them to confirm.
• ruff isn't installed locally; formatting unverified.
• No live SGLang server. The chat transport is validated against the 0.5.13 source contract, not a running server.
• nemo_gym/responses_converter.py is shared code, so this stack widens CI scope beyond sglang_model.
• The vLLM-vs-SGLang parity run requested on Add SGLang responses-API model adaptor for AR models NVIDIA-NeMo/Gym#1557 is still outstanding and needs cluster time.

Happy to reshape any of this — including holding the chat transport as a follow-up if you'd rather land NVIDIA-NeMo#1787 first.

🤖 Generated with Claude Code