feat(sglang): add token-exact SGLang model server for RL training

nemo_gym/openai_utils.py

@@ -624,3 +624,18 @@ async def create_tokenize(self, **kwargs):
 
         await self._raise_for_status(response, request_kwargs)
         return await get_response_json(response)
+
+    async def create_generate(self, **kwargs):
+        # SGLang's native generation endpoint. The public Chat-Completions
+        # response contract does not guarantee exact sampled integer token IDs,
+        # while /generate with return_logprob=True exposes selected-token IDs
+        # and logprobs. This endpoint lives at the server root, not under /v1.
+        base_url = self.base_url.removesuffix("/v1")
+        request_kwargs = dict(
+            url=f"{base_url}/generate",
+            json=kwargs,
+        )
+        response = await self._request(method="POST", **request_kwargs)
+
+        await self._raise_for_status(response, request_kwargs)
+        return await get_response_json(response)

responses_api_models/sglang_model/README.md

@@ -0,0 +1,53 @@
+# SGLang model server
+
+This Responses-API model server connects NeMo Gym to an SGLang server managed
+outside Gym. It preserves the exact prompt IDs, sampled token IDs, and sampled
+token logprobs required by token-level RL training.
+
+The adapter currently renders the chat template locally and calls SGLang's
+native `/generate` endpoint with `return_logprob=true`. Gym's public
+Chat-Completions response contract does not guarantee exact sampled integer
+token IDs, so moving this transport to `/v1/chat/completions` would currently
+lose a required training invariant. If that endpoint gains a stable
+token-ID/logprob contract, the transport can change without changing the
+session-splice or context-overflow rules below.
+
+For a multi-turn session, the adapter caches the token sequence and splices
+each prior assistant turn's exact sampled IDs into the next prompt. It never
+re-tokenizes those sampled turns. Tools and chat-template kwargs must therefore
+remain fixed for the life of a session; the adapter fails loudly if they
+change. If a prompt already fills `context_length`, the adapter returns a
+terminal response with `finish_reason="length"` instead of truncating the
+prefix. Both behaviors preserve the trainer's prefix contiguity invariant.
+
+The session cache is process-local. Run one Gym worker per model-server
+instance, or provide sticky routing that keeps every turn of a session on the
+same worker.
+
+## Configuration
+
+See `configs/sglang_model_for_training.yaml`.
+
+- `base_url`: the SGLang URL; either a bare server URL or one ending in `/v1`.
+- `model`: a tokenizer/model identifier available to the Gym server.
+- `context_length`: the SGLang server's total context limit.
+- `sglang_chat_template`: optional inline copy of the server/training template.
+- `sglang_chat_template_path`: optional path to the exact server chat template.
+- `sglang_tool_format`: `hermes` or `qwen3_coder`.
+- `trust_remote_code`: forwarded to the local tokenizer loader; defaults to
+  `false`.
+
+The leaf package pins `transformers==5.6.0`, matching NeMo RL's SGLang worker
+environment. The example config leaves `context_length` mandatory (`???`) so a
+mismatched server limit cannot be selected silently.
+
+## CPU-only tests
+
+From the Gym checkout:
+
+```bash
+uv run --extra dev pytest responses_api_models/sglang_model/tests
+```
+
+The direct tests mock the tokenizer and SGLang HTTP client; they do not load
+weights or require a GPU.

responses_api_models/sglang_model/__init__.py

@@ -0,0 +1,2 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0

responses_api_models/sglang_model/_logic.py

@@ -0,0 +1,109 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+"""Framework-free helpers for the SGLang model server."""
+
+from typing import Any, Dict, List, Tuple
+
+
+def _extract_output_ids(
+    result: Dict[str, Any],
+    meta: Dict[str, Any],
+) -> List[int] | None:
+    output_ids = meta.get("output_ids")
+    if output_ids is None:
+        output_ids = result.get("output_ids")
+    if output_ids is None:
+        return None
+    if not isinstance(output_ids, (list, tuple)):
+        raise RuntimeError(f"Malformed SGLang output_ids field: expected an array, got {type(output_ids).__name__}")
+    return [int(token_id) for token_id in output_ids]
+
+
+def _validate_selected_ids(
+    selected_ids: List[int],
+    output_ids: List[int] | None,
+) -> None:
+    if output_ids is not None and selected_ids != output_ids:
+        raise RuntimeError(
+            "SGLang returned generated-token IDs that do not match output_ids: "
+            f"selected={selected_ids!r}, output_ids={output_ids!r}"
+        )
+
+
+def extract_generated_tokens_and_logprobs(
+    result: Dict[str, Any],
+) -> Tuple[List[int], List[float]]:
+    """Return aligned generated token IDs and logprobs from ``/generate``.
+
+    SGLang releases have emitted the selected-token data as dictionaries,
+    tuples, or parallel value/index arrays. Missing or malformed data is a
+    hard error because silently returning empty arrays would invalidate the
+    training loss mask.
+    """
+    meta = result.get("meta_info") or {}
+    output_ids = _extract_output_ids(result, meta)
+    if "output_token_logprobs" in meta:
+        entries = meta["output_token_logprobs"]
+        if not isinstance(entries, (list, tuple)):
+            raise RuntimeError(
+                f"Malformed SGLang output_token_logprobs field: expected an array, got {type(entries).__name__}"
+            )
+        if not entries:
+            _validate_selected_ids([], output_ids)
+            return [], []
+        token_ids: List[int] = []
+        logprobs: List[float] = []
+        for entry in entries:
+            if isinstance(entry, dict):
+                token_id = entry.get("token_id", entry.get("id"))
+                logprob = entry.get("logprob")
+            elif isinstance(entry, (list, tuple)) and len(entry) >= 2:
+                logprob, token_id = entry[0], entry[1]
+            else:
+                raise RuntimeError(f"Malformed SGLang output_token_logprobs entry: {entry!r}")
+            if token_id is None or logprob is None:
+                raise RuntimeError(f"Malformed SGLang output_token_logprobs entry: {entry!r}")
+            token_ids.append(int(token_id))
+            logprobs.append(float(logprob))
+        _validate_selected_ids(token_ids, output_ids)
+        return token_ids, logprobs
+
+    values_present = "output_token_logprobs_val" in meta or "output_token_logprobs_val" in result
+    values = meta.get(
+        "output_token_logprobs_val",
+        result.get("output_token_logprobs_val"),
+    )
+    indexes = meta.get("output_token_logprobs_idx")
+    if indexes is None:
+        indexes = result.get("output_token_logprobs_idx")
+
+    if values_present:
+        if not isinstance(values, (list, tuple)):
+            raise RuntimeError(
+                f"Malformed SGLang output_token_logprobs_val field: expected an array, got {type(values).__name__}"
+            )
+        selected_ids = indexes if indexes is not None else output_ids
+        if not values:
+            normalized_ids = [int(token_id) for token_id in selected_ids] if selected_ids is not None else []
+            if normalized_ids:
+                raise RuntimeError(
+                    f"SGLang returned mismatched generation fields: {len(normalized_ids)} ids for 0 logprobs"
+                )
+            _validate_selected_ids(normalized_ids, output_ids)
+            return [], []
+        if not selected_ids or len(selected_ids) != len(values):
+            id_count = len(selected_ids) if selected_ids is not None else 0
+            raise RuntimeError(
+                f"SGLang returned mismatched generation fields: {id_count} ids for {len(values)} logprobs"
+            )
+        normalized_ids = [int(token_id) for token_id in selected_ids]
+        _validate_selected_ids(normalized_ids, output_ids)
+        return normalized_ids, [float(logprob) for logprob in values]
+
+    result_keys = sorted(result)
+    meta_keys = sorted(meta)
+    raise RuntimeError(
+        "SGLang /generate returned no generated-token logprobs "
+        f"(result keys={result_keys}, meta_info keys={meta_keys}). "
+        "Ensure return_logprob=true is supported and honored."
+    )