Skip to content

Inference: Bring chat completions API inline with vllm/official openAI spec #5276

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
sidsingh-nvidia merged 5 commits into from
Jul 14, 2026
Merged

Inference: Bring chat completions API inline with vllm/official openAI spec #5276

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
9773d0e
delegate reasoning token retention to the chat template
sidsingh-nvidia Jun 10, 2026
c321ab6
trim chat completions content
sidsingh-nvidia Jun 25, 2026
6042300
Merge branch 'main' into siddharth/fix-think-token-retention-multi-turn
sidsingh-nvidia Jun 25, 2026
26d5843
address Robert's feedback
sidsingh-nvidia Jul 13, 2026
8f9657c
flag for returning raw text optionally and lint
sidsingh-nvidia Jul 13, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 0 additions & 4 deletions megatron/core/inference/data_parallel_inference_coordinator.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -586,10 +586,6 @@ def detokenize(self, finished_request):
finished_request (dict): The serialized merged request containing the
generated tokens to be detokenized. It is modified in place.
"""
if finished_request["prompt"] is None:
finished_request["prompt"] = TextGenerationController.detokenize(
self.tokenizer, finished_request["prompt_tokens"][1], remove_EOD=False
)
detokenize_stop_sequence = (finished_request.get("sampling_params", {}) or {}).get(
"detokenize_stop_sequence", False
)
Expand Down
48 changes: 22 additions & 26 deletions ...re/inference/text_generation_server/dynamic_text_gen_server/endpoints/chat_completions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -323,22 +323,6 @@ def _sanitize_tools_for_template(tools):
return sanitized


def _reconstruct_reasoning_content(messages: list[dict]) -> list[dict]:
"""Reconstruct <think> tags from reasoning_content fields on assistant messages.

For parity with vLLM, assistant messages may carry reasoning in the reasoning_content field.
Before applying the chat template, we must inline those tags back into content.
"""
for message in messages:
if message.get("role") != "assistant":
continue
reasoning_content = message.pop("reasoning_content", None)
if reasoning_content is not None:
content = message.get("content") or ""
message["content"] = f"<think>{reasoning_content}</think>{content}"
return messages


def _replace_prefix_tokens(
eos_token_id,
previous_turn_token_ids,
Expand Down Expand Up @@ -462,7 +446,6 @@ async def chat_completions():
if not isinstance(messages, list):
return Response("'messages' must be a list", status=400)
template_messages = _sanitize_messages_for_template(messages)
template_messages = _reconstruct_reasoning_content(template_messages)
template_tools = _sanitize_tools_for_template(tools)

try:
Expand Down Expand Up @@ -592,6 +575,7 @@ async def chat_completions():
prompt_tokens = [tokenizer.bos] + prompt_tokens

max_tokens = req.get("max_completion_tokens", None) or req.get("max_tokens", None)
ignore_eos = bool(req.get("ignore_eos", False))

sampling_params = SamplingParams(
temperature=temperature,
Expand All @@ -602,6 +586,7 @@ async def chat_completions():
num_tokens_to_generate=(int(max_tokens) if max_tokens is not None else None),
skip_prompt_log_probs=skip_prompt_log_probs,
add_BOS=add_BOS,
termination_id=-1 if ignore_eos else None,
)
except ValueError as e:
return Response(f"Invalid sampling parameter: {e}", status=400)
Expand Down Expand Up @@ -663,6 +648,14 @@ async def chat_completions():
total_completion_tokens = 0
prompt_tokens_counts = []

prevent_retokenization = req.get("prevent_retokenization", True)

@ArEsKay3 ArEsKay3 Jun 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This isn't really true. MRL uses prompt/generation token ids but doesn't always want to prevent retokenization. Getting the token ids back is a separate feature than this.

@ArEsKay3 ArEsKay3 Jun 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I think it's probably better to gate this off of a new switch?

@sidsingh-nvidia sidsingh-nvidia Jun 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Will do.

@sidsingh-nvidia sidsingh-nvidia Jun 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We can have a flag called "return_tokenized_data", defaulting to False to match out-of-the-box vLLM.

  1. Do we wanna set this to True in megatron-RL?
  2. If prevent_retokenization is True, then do we always return tokenized data regardless of the value of the new flag?

@sidsingh-nvidia sidsingh-nvidia Jul 13, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Addressed below. I have added a "return_tokenized_data" flag/

# return_tokenized_data controls whether prompt/generation token ids are
# included in the response. It is independent of prevent_retokenization
# (a client may want token ids without prevent_retokenization, or vice versa),
# but prevent_retokenization implicitly requires token ids so the client
# can echo them back next turn.
return_tokenized_data = req.get("return_tokenized_data", False) or prevent_retokenization
return_raw_text = req.get("return_raw_text", False)
request_idx = 0
for result_item in batch_results:
result = unwrap_serialized_tensors(result_item)
Expand Down Expand Up @@ -736,9 +729,14 @@ async def chat_completions():
if "reasoning" in metadata:
message["reasoning_content"] = metadata["reasoning"]

# Replicate data in the message field for compatibility.
message["prompt_token_ids"] = result["prompt_tokens"]
message["generation_token_ids"] = result["generated_tokens"]
if return_tokenized_data:
message["prompt_token_ids"] = result["prompt_tokens"]
message["generation_token_ids"] = result["generated_tokens"]
if return_raw_text:
prompt_str = tokenizer.detokenize(result["prompt_tokens"])
message["raw_text"] = prompt_str + text_output
# Small RL/debug scalars (a few bytes each); harmless to keep for
# NeMo-RL compatibility.
message["generation_log_probs"] = result.get("generated_log_probs", [])
Comment on lines +738 to 740

@claude claude Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION Naming] The comment "Small RL/debug scalars (a few bytes each)" accurately describes policy_epoch, kv_cache_epoch, and num_evictions, but it is placed on a block whose first member is generation_log_probs — a per-generated-token float list, not a scalar. For a long completion this is one of the larger fields in the payload, which contradicts both the "a few bytes each" wording and the PR's stated goal of cutting payload size.

Why it matters: a wrong comment is worse than none — a future reader trimming the payload may skip generation_log_probs believing it is negligible.

Suggested fix: move the comment below generation_log_probs so it only covers the three scalar fields, and note that generation_log_probs is always included because RL clients (e.g. megatron/rl/inference/megatron.py) read choice.message.generation_log_probs. If it should also be gated behind an opt-in flag like the token ids, consider doing so for consistency.

Comment on lines +738 to 740

@claude claude Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[SUGGESTION Naming] The comment "Small RL/debug scalars (a few bytes each)" accurately describes policy_epoch, kv_cache_epoch, and num_evictions, but it is applied to a block whose first member is generation_log_probs — a per-generated-token float list, not a scalar. For a long completion this is one of the larger fields in the payload, which directly contradicts the "a few bytes each" characterization and the PR's stated goal of cutting payload size.

Why it matters: a wrong comment is worse than none — a future reader trimming the payload may skip generation_log_probs believing it is negligible.

Suggested fix: move the comment below generation_log_probs so it only covers the three scalar fields, and note separately that generation_log_probs is always included (it is required by RL clients such as megatron/rl/inference/megatron.py, which reads choice.message.generation_log_probs). If it should also be gated behind an opt-in flag like the token ids, consider doing so for consistency.

message["policy_epoch"] = result["policy_epoch"]
message["kv_cache_epoch"] = result["kv_cache_epoch"]
Expand All @@ -759,15 +757,13 @@ async def chat_completions():
else:
finish_reason = "stop"

# Choice-level prompt/generation_token_ids, generation_log_probs and
# raw_text were duplicates of message-level data (or reconstructable);
# dropped to match vLLM's response shape and cut payload size.
choice_data = {
"index": request_idx,
"message": message,
"prompt_token_ids": result["prompt_tokens"],

@ArEsKay3 ArEsKay3 Jun 25, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Removing this from the choice data and only having it in the message data is probably good to remove reduncancy, but this will require a corresponding change in MRL without NemoGym.
Currently, we are using this field rather than the message object.
https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/rl/inference/megatron.py#L77

@sidsingh-nvidia sidsingh-nvidia Jun 26, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I can make this change in this PR.

@sidsingh-nvidia sidsingh-nvidia Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done

@ArEsKay3 ArEsKay3 Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

👍

"generation_token_ids": result["generated_tokens"],
"generation_log_probs": result.get("generated_log_probs", []),
"raw_text": result["prompt"] + result["generated_text"],
# 'logprobs' in chat API is an object containing 'content'
# "logprobs": {"content": logprobs_content} if logprobs_content else None,
"logprobs": {"content": logprobs_content} if return_log_probs else None,
"finish_reason": finish_reason,
}
Comment on lines +760 to 769

@claude claude Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] This removes prompt_token_ids, generation_token_ids, generation_log_probs, and raw_text from the choice object. The in-repo consumer (megatron/rl/inference/megatron.py) is updated in this PR to read them from choice.message.* instead, so the internal path is fine. But this is a response-shape change on a public OpenAI-compatible endpoint: any external client that read these at the choice level (e.g. a NeMo-RL client pinned to the previous shape) will silently get None/AttributeError after this change, with no deprecation window.

Why it matters: silent breakage of an external client contract is hard to diagnose from the server side.

Suggestion: confirm NeMo-RL (and any other downstream) reads the message-level fields, and if this endpoint is externally consumed, call out the shape change in the PR description / release notes. No code change required if downstreams are already aligned.

@sidsingh-nvidia sidsingh-nvidia Jul 14, 2026

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is this deviating from the official openAI spec?

@ArEsKay3 ArEsKay3 Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No, these are not part of the spec.

Comment on lines +760 to 769

@claude claude Bot Jul 14, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[IMPORTANT Compatibility] This removes prompt_token_ids, generation_token_ids, generation_log_probs, and raw_text from the choice object. The in-repo consumer (megatron/rl/inference/megatron.py) is updated in this PR to read them from choice.message.* instead, so the internal path is fine. However, this is a response-shape change on a public OpenAI-compatible endpoint: any external client that read these at the choice level (e.g. a NeMo-RL client pinned to the previous shape) will silently get None/AttributeError after this change, with no deprecation window.

Why it matters: silent breakage of an external client contract is hard to diagnose from the server side.

Suggestion: confirm the NeMo-RL (and any other downstream) client reads the message-level fields, and if this endpoint is externally consumed, call out the shape change in the PR description / release notes so consumers can migrate. No code change required if downstreams are already aligned.

Expand All @@ -782,7 +778,7 @@ async def chat_completions():
]

choices.append(choice_data)
if choice_data["generation_log_probs"] is None:
if result.get("generated_log_probs") is None:
logger.warning(
"Generation log probs is None for request:\n%s",
json.dumps(_redact_token_id_lists_for_logging(result), indent=4),
Expand Down
14 changes: 10 additions & 4 deletions megatron/rl/inference/megatron.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,12 @@ async def base_generate(self, request: InferenceRequest) -> InferenceResponse:
extra_body={
"skip_prompt_log_probs": True,
"add_BOS": (not args.rl_skip_bos_token and tokenizer.bos is not None),
# TODO: These are non-standard fields that add significant memory overheads to the
# chat completions payload. return_raw_text also wastes a lot of CPU cycles
# detokenizing prompt tokens, especially expensive for long prompts in agentic RL.
# Set to False if not needed in MRL.
"return_tokenized_data": True,
"return_raw_text": True,
},
)

Expand All @@ -73,11 +79,11 @@ async def base_generate(self, request: InferenceRequest) -> InferenceResponse:
return InferenceResponse(
# TODO: Handle tool calls and reasoning in LLMChatMessage
response=LLMChatMessage(**choice.message.model_dump(include={'role', 'content'})),
raw_text=choice.raw_text,
token_ids=choice.prompt_token_ids + choice.generation_token_ids,
logprobs=choice.generation_log_probs,
raw_text=choice.message.raw_text,
token_ids=choice.message.prompt_token_ids + choice.message.generation_token_ids,
logprobs=choice.message.generation_log_probs,
finish_reason=choice.finish_reason,
prompt_length=len(choice.prompt_token_ids),
prompt_length=len(choice.message.prompt_token_ids),
policy_epoch=choice.message.policy_epoch,
kv_cache_epoch=choice.message.kv_cache_epoch,
num_evictions=choice.message.num_evictions,
Expand Down
Loading