Skip to content

Add end-to-end sampling controls to the LLM server worker protocol - #21561

Merged
metascroy merged 2 commits into
mainfrom
support-sampling-params
Aug 4, 2026
Merged

Add end-to-end sampling controls to the LLM server worker protocol#21561
metascroy merged 2 commits into
mainfrom
support-sampling-params

Conversation

@metascroy

Copy link
Copy Markdown
Contributor

Summary

Add end-to-end sampling controls to the LLM server worker protocol.

  • Plumb top_p, top_k, and seed from chat completion requests through the Python runtime and JSONL worker protocol.
  • Extend the C++ worker sampling configuration while preserving existing defaults and greedy behavior.
  • Validate sampling parameters at the API boundary and document their supported ranges and semantics.
  • Add coverage for request validation, worker serialization, session propagation, and C++ protocol handling.

Test Plan

  • Ran the focused Python LLM server test suite: 78 tests passed.
  • Built the C++ worker integration successfully.
  • Verified deterministic output for repeated seeded requests, different output for different or omitted seeds, and expected behavior for top_p and top_k.
  • Verified invalid sampling values return HTTP 400 responses.

@pytorch-bot

pytorch-bot Bot commented Aug 3, 2026

Copy link
Copy Markdown

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/21561

Note: Links to docs will display an error until the docs builds have been completed.

⏳ No Failures, 116 Pending

As of commit 1eca171 with merge base 4b4df96 (image):
💚 Looks good so far! There are no failures yet. 💚

This comment was automatically generated by Dr. CI and updates every 15 minutes.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Aug 3, 2026
@metascroy

Copy link
Copy Markdown
Contributor Author

@claude review this code

@github-actions

github-actions Bot commented Aug 4, 2026

Copy link
Copy Markdown

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

@claude

claude Bot commented Aug 4, 2026

Copy link
Copy Markdown

Claude finished @metascroy's task in 3m 2s —— View job


Code Review

Reviewed the full diff (git show HEAD) plus SamplingConfig in extension/llm/runner/llm_session.h. This is a clean, well-tested change — plumbing is consistent across the Python request → GenerationOptions_WorkerRequest → JSONL → C++ SamplingConfig path, and each layer has coverage. A few things worth addressing:

🔴 Duplicated comment (leftover from edit)

examples/llm_server/cpp/worker_loop.h:132-133

  // Stop strings
  // Stop strings (the request's `stop` sequences): terminate at the token

Line 132 is a stray duplicate — the real comment starts on 133. Drop line 132.

Fix this →

🟡 seed validation diverges between the Python API and the C++ worker

  • Python (serving_chat.py:373): 0 < req.seed <= 2**63 - 1 — rejects an explicit seed=0.
  • C++ (worker_loop.h:129): only seed_value < 0 throws — accepts 0.

This is defensible (Python treats 0 as the "unset" sentinel it forwards when seed is omitted, and SamplingConfig::seed == 0 means "unset"), and the C++ side must accept 0 because Python always forwards 0 for an omitted seed. But the two layers tell a slightly different story, and a caller talking directly to the worker gets different semantics than one going through the HTTP API. Worth a one-line comment on the C++ side noting 0 == unset so the asymmetry is intentional and documented, matching the README's "omitted seed uses the worker's unset/random value."

Also note the C++ upper bound isn't guarded: req.value("seed", ...) into an int64_t will overflow / throw for a JSON seed ≥ 2⁶³ when hitting the worker directly (the Python layer catches this at 2**63, C++ doesn't). Low severity since the tested path is Python-fronted, but the C++ tests only cover seed = -1.

🟡 seed/top_p/top_k are silently no-ops under greedy decoding

Default temperature is 0.0 (greedy). With greedy decoding, top_p/top_k/seed have no effect, so a client sending seed=42, temperature=0 gets deterministic-but-not-seeded output and no error. That's arguably fine, but since the PR advertises "deterministic output for repeated seeded requests," it's worth confirming the worker's sampler actually honors seed only when temperature > 0, and whether the README should note that these knobs require temperature > 0.

🟢 Minor

  • examples/llm_server/python/tests/test_worker_client.py:16-17 — the blank line between import pytest (third-party) and the first-party executorch... import was removed, merging the import groups. lintrunner (isort/ufmt) may flag this; worth running lintrunner -a before merge.
  • The validation-before-prefill ordering in C++ is correct and the prefill_calls == 0 assertion in test_invalid_sampling_config_rejected nicely guards it. 👍
  • std::isfinite(top_p_value) on the parsed double correctly rejects NaN/Inf; matches the Python math.isfinite check. Ranges (0, 1] for top_p and [0, 2³¹−1] for top_k are consistent across both layers. 👍

Summary

No blocking bugs. The duplicated comment is a trivial cleanup; the seed semantics divergence and greedy-decoding no-op behavior are worth a short comment / README note but are intentional design choices. Test coverage is thorough across all layers.
· support-sampling-params

@metascroy
metascroy merged commit 454f979 into main Aug 4, 2026
195 of 197 checks passed
@metascroy
metascroy deleted the support-sampling-params branch August 4, 2026 21:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants