Skip to content

[SkyRL][train] Support prompt_logprobs in /sample in the new inference stack#1417

Open
nithinvc wants to merge 6 commits intoNovaSky-AI:mainfrom
nithinvc:nithinc/sample-prompt-logprobs
Open

[SkyRL][train] Support prompt_logprobs in /sample in the new inference stack#1417
nithinvc wants to merge 6 commits intoNovaSky-AI:mainfrom
nithinvc:nithinc/sample-prompt-logprobs

Conversation

@nithinvc
Copy link
Copy Markdown
Contributor

@nithinvc nithinvc commented Mar 31, 2026

Summary

This PR addresses #1416 and adds prompt log probabilities to the new /sample endpoint in the new remote inference client.

  • Add prompt_logprobs and topk_prompt_logprobs support to RemoteInferenceClient.sample(), wiring through to vLLM's /inference/v1/generate endpoint
  • Accepts include_prompt_logprobs (bool) and topk_prompt_logprobs (int) from the request body, with fallback to prompt_logprobs for Tinker SampleRequest compatibility

Test plan

  • CPU: test_sample — verifies include_prompt_logprobs=True returns per-token logprobs with correct values
  • CPU: test_sample_n2 — verifies prompt logprobs work with n=2
  • CPU: test_sample_topk_prompt_logprobs — verifies topk_prompt_logprobs=2 returns both fields with correct (token_id, logprob) values
  • GPU: test_client_sample — verifies vLLM returns prompt logprobs (non-positive floats)
  • GPU: test_client_sample_multiple — verifies prompt logprobs with n=3
  • GPU: test_client_sample_topk_prompt_logprobs — verifies that the vLLM top-k tuples contain the actual prompt token

Open with Devin

@nithinvc nithinvc marked this pull request as ready for review March 31, 2026 01:25
gemini-code-assist[bot]

This comment was marked as resolved.

Copy link
Copy Markdown
Contributor

@devin-ai-integration devin-ai-integration bot left a comment

Choose a reason for hiding this comment

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

✅ Devin Review: No Issues Found

Devin Review analyzed this PR and found no potential bugs to report.

View in Devin Review to see 4 additional findings.

Open in Devin Review

@nithinvc nithinvc marked this pull request as draft March 31, 2026 02:03
@nithinvc nithinvc marked this pull request as ready for review March 31, 2026 02:58
gemini-code-assist[bot]

This comment was marked as resolved.

…client.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
devin-ai-integration[bot]

This comment was marked as resolved.

@SumanthRH SumanthRH self-assigned this Mar 31, 2026
…client.py

Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants