Refactor of vLLM bridge #1914
Open
Refactor of vLLM bridge #1914
Conversation
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
jreiffers
reviewed
Jan 26, 2026
| """ | ||
|
|
||
| def pad_left(seq: List[int], length: int, pad_value: int) -> List[int]: | ||
| seq = seq[:length] # Truncate if too long |
Member
There was a problem hiding this comment.
Hmm, silently truncating feels wrong.
Contributor
Author
There was a problem hiding this comment.
Updated to not silently truncate anymore
jreiffers
reviewed
Jan 26, 2026
| """Get all generated token ID sequences.""" | ||
| return [c.token_ids for c in self.completions] | ||
|
|
||
| def to_arrays( |
Member
There was a problem hiding this comment.
This seems to be unused? Am I missing something?
Contributor
Author
There was a problem hiding this comment.
I have a simpler debug example of an RL loop that doesn't use Tunix that uses it, I just pushed that example as well
jreiffers
reviewed
Jan 26, 2026
| - Dict[str, jax.Array]: Direct flattened params | ||
| - flax.nnx.State: Flax state object | ||
| - flax.nnx.Module: Flax module (state extracted automatically) | ||
| block: If True, wait for transfer completion (always True currently). |
Member
There was a problem hiding this comment.
Why is this flag there if it doesn't do anything?
jreiffers
reviewed
Jan 26, 2026
| assert len(original) <= length, f"Sequence too long: {len(original)} > {length}" | ||
| return original + [pad_value] * (length - len(original)) | ||
|
|
||
| for i, completion in enumerate(output.completions): |
jreiffers
reviewed
Jan 26, 2026
| input_tokens = [] | ||
| output_tokens = [] | ||
|
|
||
| def pad_to_left(original: List[int], length: int, pad_value: int) -> List[int]: |
Member
There was a problem hiding this comment.
There's another implementation of this in api/types.py.
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 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.
The purpose of the PR is to refactor the existing JAX <> vLLM bridge to expose and interface that's more easily usable by user who want to avoid using Tunix, the main change is the addition of a framework agnostic interface -
VLLMRolloutEnginethat an external RL framework can inherit / wrap around to access as the entry point to the functionality of the bridge