feat(RL): add vLLM Tokens-in-Tokens-Out and RL related response support - #9651
Merged
Conversation
biswapanda
changed the title
biswapanda
changed the title
dynamo-ops
reviewed
May 16, 2026
Contributor
Author
|
@coderabbitai review |
Contributor
✅ Actions performedReview triggered.
|
Contributor
WalkthroughThis PR implements RL-mode support with conditional logprobs-mode switching and nvext-aware token handling. It tracks explicit ChangesvLLM RL mode and nvext token handling
🎯 3 (Moderate) | ⏱️ ~25 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (2)
components/src/dynamo/vllm/tests/test_vllm_unit.py (1)
361-372: ⚡ Quick winAdd coverage for
--logprobs-mode=<value>flag form.
_arg_was_provided()handles both split and=forms; this test currently exercises only split args. Covering both prevents regressions in explicit-flag tracking.✅ Suggested test tweak
-def test_logprobs_mode_flag_is_tracked(mock_vllm_cli): - mock_vllm_cli( - "--model", - "Qwen/Qwen3-0.6B", - "--logprobs-mode", - "raw_logprobs", - ) - - config = parse_args() - - assert config.logprobs_mode_explicitly_set is True +@pytest.mark.parametrize( + "logprobs_args", + [ + ("--logprobs-mode", "raw_logprobs"), + ("--logprobs-mode=raw_logprobs",), + ], +) +def test_logprobs_mode_flag_is_tracked(mock_vllm_cli, logprobs_args): + mock_vllm_cli( + "--model", + "Qwen/Qwen3-0.6B", + *logprobs_args, + ) + + config = parse_args() + + assert config.logprobs_mode_explicitly_set is True🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@components/src/dynamo/vllm/tests/test_vllm_unit.py` around lines 361 - 372, The test test_logprobs_mode_flag_is_tracked currently only exercises the split-argument form; add a complementary test (or extend this one) that calls mock_vllm_cli with the equals form (e.g., "--logprobs-mode=raw_logprobs"), then call parse_args() and assert config.logprobs_mode_explicitly_set is True to ensure _arg_was_provided() correctly recognizes the '=' form as explicitly set.components/src/dynamo/vllm/handlers.py (1)
398-406: ⚡ Quick winAvoid quadratic flattening in
_flatten_logprobs.
pending.pop(0)pluspending[0:0] = itemmakes this O(n²). On long RL/TITO logprob payloads this helper can become a hot path for no real benefit.♻️ Suggested change
- pending = list(log_probs) + pending = list(reversed(log_probs)) while pending: - item = pending.pop(0) + item = pending.pop() if isinstance(item, (int, float)): out.append(float(item)) elif isinstance(item, list): - pending[0:0] = item + pending.extend(reversed(item)) elif isinstance(item, dict) and "logprob" in item: try: out.append(float(item["logprob"])) except (TypeError, ValueError): continue🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@components/src/dynamo/vllm/handlers.py` around lines 398 - 406, Replace the quadratic list-front-insert pattern in _flatten_logprobs: instead of using pending as a list with pending.pop(0) and pending[0:0] = item (which causes O(n^2)), make pending a collections.deque (initialized from log_probs), use pending.popleft() instead of pop(0), and when encountering a list item use pending.extendleft(reversed(item)) to prepend the sublist efficiently; keep the same handling for numeric types and dicts with "logprob" and append floats to out as before.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@components/src/dynamo/vllm/handlers.py`:
- Around line 2616-2623: The current code sets request_prompt_token_ids from
request.get("token_ids"), which is incorrect for the pre-rendered multimodal
path; change the assignment so that when want_engine_data is true you prefer the
built prompt's token ids (the expanded sequence's prompt["prompt_token_ids"])
and only fall back to request.get("token_ids") otherwise; update the logic that
populates request_prompt_token_ids (the variable used to populate
engine_data.prompt_token_ids) so it reads from prompt["prompt_token_ids"] when
present to ensure engine_data echoes the actual prompt the engine consumed.
In `@components/src/dynamo/vllm/tests/test_vllm_tito_parity.py`:
- Around line 13-22: The pytest module-level marker list pytestmark (which
already includes pytest.mark.vllm) is missing the required single component
marker; update the pytestmark list in test_vllm_tito_parity.py to include
exactly one component marker (choose one of pytest.mark.multimodal,
pytest.mark.router, pytest.mark.kvbm, or pytest.mark.core) alongside the
existing markers so the module has a framework marker (vllm) and exactly one
component marker.
---
Nitpick comments:
In `@components/src/dynamo/vllm/handlers.py`:
- Around line 398-406: Replace the quadratic list-front-insert pattern in
_flatten_logprobs: instead of using pending as a list with pending.pop(0) and
pending[0:0] = item (which causes O(n^2)), make pending a collections.deque
(initialized from log_probs), use pending.popleft() instead of pop(0), and when
encountering a list item use pending.extendleft(reversed(item)) to prepend the
sublist efficiently; keep the same handling for numeric types and dicts with
"logprob" and append floats to out as before.
In `@components/src/dynamo/vllm/tests/test_vllm_unit.py`:
- Around line 361-372: The test test_logprobs_mode_flag_is_tracked currently
only exercises the split-argument form; add a complementary test (or extend this
one) that calls mock_vllm_cli with the equals form (e.g.,
"--logprobs-mode=raw_logprobs"), then call parse_args() and assert
config.logprobs_mode_explicitly_set is True to ensure _arg_was_provided()
correctly recognizes the '=' form as explicitly set.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Pro
Run ID: cfb6dcd4-cd82-48c4-aed4-79646146e466
📒 Files selected for processing (6)
components/src/dynamo/vllm/args.pycomponents/src/dynamo/vllm/backend_args.pycomponents/src/dynamo/vllm/handlers.pycomponents/src/dynamo/vllm/main.pycomponents/src/dynamo/vllm/tests/test_vllm_tito_parity.pycomponents/src/dynamo/vllm/tests/test_vllm_unit.py
biswapanda
changed the title
biswapanda
changed the title
dynamo-ops
approved these changes
May 16, 2026
biswapanda
force-pushed
the
bis/nvext-tito-vllm
branch
from
May 16, 2026 01:02
cac71f2 to
8cec894
Compare
dynamo-ops
reviewed
May 16, 2026
biswapanda
force-pushed
the
bis/nvext-tito-rl
branch
from
May 16, 2026 01:12
565fd7c to
2256c2a
Compare
biswapanda
force-pushed
the
bis/nvext-tito-vllm
branch
from
May 16, 2026 01:13
8cec894 to
a836fc9
Compare
dynamo-ops
reviewed
May 16, 2026
biswapanda
force-pushed
the
bis/nvext-tito-vllm
branch
from
May 16, 2026 01:57
fae4010 to
b9eb3bd
Compare
Emit dtype alongside data/shape so the consumer decodes the raw base85 bytes with the correct element type instead of assuming int32 (KrishnanPrash, codex).
detokenize is already forced False unconditionally at the end of the function; the earlier duplicate assignment was dead.
build_sampling_params forces detokenize=False, so vLLM never reads skip_special_tokens; Dynamo detokenizes in the Rust backend, which reads it from the request output_options (backend.rs). Replace the no-op assignment with a NOTE and update the tests to assert it is not forwarded (KrishnanPrash).
Capture the raw routed_experts per output index during streaming and base85- encode it once on the finish chunk, instead of serializing on every chunk and discarding all but the last. Non-final chunks now only do a getattr.
Skip a non-int prompt-logprob token-id key instead of aborting the whole prompt_logprobs payload (kthui), and drop a malformed completion token id instead of killing the generation stream, in _accumulate_engine_data.
The two metadata/engine_data extra_fields tests ran with enable_rl=False, which short-circuits the RL branch so they validated nothing about the gate. Run them with enable_rl=True so they confirm metadata-only requests are not treated as token-in even when RL is active.
Lock the -inf/nan -> finite sentinel behavior in _finite_logprob, _flatten_logprobs (incl. bool drop), and _serialize_prompt_logprobs, and assert the sentinel is JSON-safe so the serde_json transport never nulls it.
biswapanda
force-pushed
the
bis/nvext-tito-vllm
branch
from
June 6, 2026 07:37
7a2a3cd to
68f5c1c
Compare
Contributor
Author
|
rebase on top of latest main |
dict(out.get('disaggregated_params') or {}) tripped mypy's arg-type check
(the value union includes list/int). Narrow with isinstance before dict(),
which is also runtime-safer. Fixes the dynamo-runtime/mypy CI failure.
prefill_result.get('disaggregated_params', {}) returns None (not {}) when the
key is present-but-None, which prefill produces on its error path and when
_build_disaggregated_params returns None for empty params. The chained .get()
then raised AttributeError. Use 'or {}' so missing / None / empty all degrade
to an empty mapping.
VllmLLMEngine.generate (unified path) had the same prefill_result
.get('disaggregated_params', {}).get(...) trap as the legacy handler: a
present-but-None value -> AttributeError. Use 'or {}' so it falls through to
the existing kv_params ValueError instead.
Absorbed from the RL worker-admin stack so the binding lives in the base PR. Workers use it for their RL request-plane route descriptor instead of deriving the system URL from static env vars. Uses ip_resolver::local_ip_for_advertise (renamed from get_local_ip_for_advertise on current main).
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
7 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.
Summary
nvext.cache_salton both decode and disaggregated prefill promptsnvext.engine_data, completion token IDs/logprobs, prompt logprobs, and related TITO parity fieldsStack
bis/nvext-tito-rl) with the Rustlib/llmprotocol/preprocessor plumbingcomponents/src/dynamo/vllm/*Where To Review
components/src/dynamo/vllm/handlers.pycomponents/src/dynamo/vllm/tests/test_vllm_tito_parity.pycomponents/src/dynamo/vllm/args.pycomponents/src/dynamo/vllm/backend_args.pycomponents/src/dynamo/vllm/main.pyLocal Tests
python -m py_compile components/src/dynamo/vllm/handlers.py components/src/dynamo/vllm/tests/test_vllm_tito_parity.pyPYTHONPATH=components/src uv run --no-project --with pytest --with pytest-benchmark python -m pytest -c /dev/null components/src/dynamo/vllm/tests/test_vllm_tito_parity.py -qgit diff --checkNotes
After #9649 lands, this PR can be retargeted to
mainor rebased so the diff remains limited to the vLLM files.supersedes #9382
Summary by CodeRabbit
Release Notes
New Features
--enable-rlconfiguration option for reinforcement learning-style defaults.bad_words_token_idsandskip_special_tokenssampling parameters.Tests