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

## 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.

Author: metascroy

examples/llm_server/README.md

@@ -26,14 +26,17 @@ See `python/README.md` to run it.
 Status: experimental, reliability-first and deliberately narrow. Implemented:
 `/health`, `/v1/models`, `/v1/chat/completions` (streaming + non-streaming),
 Hugging Face chat templates (`--hf-tokenizer`), `temperature` / `max_tokens` /
-`max_completion_tokens` / `stop`, Hermes tool calling by default
+`max_completion_tokens` / `stop`, sampling controls `top_p` / `top_k` / `seed`
+(these only affect output when `temperature > 0`; under the default greedy
+decoding `temperature = 0` they are accepted but have no effect, and an omitted
+`seed` uses the worker's unset/random value), Hermes tool calling by default
 (`<tool_call>...</tool_call>` JSON, complete calls only; model-specific launchers
 may select the Qwen XML format) with `tool_choice="none"`,
 structured API errors, and best-effort cancellation. One worker process with
 serialized execution; a worker can host isolated sessions on one weight load when its engine reports
 capacity > 1 (with warm append-only resume across turns). KV/prefix state lives inside the
-worker/session, not the control plane. Unsupported params (including `top_p`,
-`seed`, `n>1`, `reasoning_effort`, penalties, `logit_bias`, `response_format`,
+worker/session, not the control plane. Unsupported params (including
+`n>1`, `reasoning_effort`, penalties, `logit_bias`, `response_format`,
 `logprobs`, and `tool_choice="required"`) are rejected with a structured 400
 rather than silently ignored. See `python/README.md` to run it and
 `spec/README.md` for the exact contract.
@@ -85,8 +88,10 @@ Supported contract for pi:
 - `tool_choice`: only `"auto"`, `"none"`, or unset.
 - Rejected with a structured 400 (`unsupported_parameter`), not silently
   ignored: `tool_choice="required"` or specific-function forcing,
-  `response_format` JSON/constrained output, `logprobs`, `top_p` other than
-  `1.0`, and `seed`.
+  `response_format` JSON/constrained output, and `logprobs`.
+- `top_p`, `top_k`, and `seed` are supported, but only take effect when
+  `temperature > 0`; the default greedy decoding (`temperature = 0`) ignores
+  them, and an omitted `seed` uses the worker's unset/random value.
 
 Reliability guidance:
 

examples/llm_server/cpp/test_worker_loop.cpp

@@ -18,6 +18,7 @@
 #include <cstdio>
 #include <iostream>
 #include <iterator>
+#include <limits>
 #include <sstream>
 #include <string>
 #include <unordered_map>
@@ -64,13 +65,18 @@ class FakeSession : public LLMSession {
   int fail_prefill_on = -1; // 0-based call index to fail (-1 = never)
   int decode_calls = 0;
   int fail_decode_on = -1;
+  std::vector<SamplingConfig> prefill_sampling;
+  std::vector<SamplingConfig> decode_sampling;
   int reset_calls = 0;
   bool fail_reset = false;
 
   ETError prefill_tokens(
       const std::vector<uint64_t>& tokens,
-      const SamplingConfig* /*initial_sampling*/ = nullptr) override {
+      const SamplingConfig* initial_sampling = nullptr) override {
     prefill_sizes.push_back(tokens.size());
+    if (initial_sampling != nullptr) {
+      prefill_sampling.push_back(*initial_sampling);
+    }
     prefill_batches.push_back(tokens);
     if (prefill_calls++ == fail_prefill_on) {
       return ETError::Internal; // failed AFTER (notionally) mutating state
@@ -79,7 +85,8 @@ class FakeSession : public LLMSession {
     return ETError::Ok;
   }
 
-  ETResult<DecodeResult> decode_one(const SamplingConfig& /*s*/) override {
+  ETResult<DecodeResult> decode_one(const SamplingConfig& sampling) override {
+    decode_sampling.push_back(sampling);
     if (decode_calls++ == fail_decode_on) {
       return ETError::Internal;
     }
@@ -219,6 +226,67 @@ nlohmann::json idsReq(std::vector<uint64_t> ids, int64_t max_new = 8) {
   return {{"max_new_tokens", max_new}, {"prompt_segments", {{{"ids", ids}}}}};
 }
 
+bool sameSampling(
+    const SamplingConfig& sampling,
+    float temperature,
+    float top_p,
+    int32_t top_k,
+    uint64_t seed) {
+  return sampling.temperature == temperature && sampling.top_p == top_p &&
+      sampling.top_k == top_k && sampling.seed == seed;
+}
+
+void test_sampling_config_forwarded() {
+  auto st = makeState();
+  fake(st).steps = {{10, "a", false, false}, {0, "", true, true}};
+  auto req = idsReq({1, 2, 3});
+  req["temperature"] = 0.7;
+  req["top_p"] = 0.8;
+  req["top_k"] = 32;
+  req["seed"] = 123;
+  run(st, /*warm=*/false, req);
+
+  check(
+      "sampling: explicit config reaches prefill",
+      fake(st).prefill_sampling.size() == 1 &&
+          sameSampling(fake(st).prefill_sampling[0], 0.7f, 0.8f, 32, 123));
+  check(
+      "sampling: explicit config reaches every decode",
+      fake(st).decode_sampling.size() == 2 &&
+          std::all_of(
+              fake(st).decode_sampling.begin(),
+              fake(st).decode_sampling.end(),
+              [](const SamplingConfig& sampling) {
+                return sameSampling(sampling, 0.7f, 0.8f, 32, 123);
+              }));
+
+  auto defaults = makeState();
+  fake(defaults).steps = {{0, "", true, true}};
+  run(defaults, /*warm=*/false, idsReq({1}));
+  check(
+      "sampling: omitted fields use worker defaults",
+      fake(defaults).prefill_sampling.size() == 1 &&
+          sameSampling(fake(defaults).prefill_sampling[0], 0.0f, 1.0f, 0, 0));
+}
+
+void test_invalid_sampling_config_rejected() {
+  for (auto invalid : std::vector<nlohmann::json>{
+           {{"top_p", 0.0}},
+           {{"top_k", -1}},
+           {{"top_k",
+             static_cast<int64_t>(std::numeric_limits<int32_t>::max()) + 1}},
+           {{"seed", -1}},
+       }) {
+    auto st = makeState();
+    auto req = idsReq({1});
+    req.update(invalid);
+    const auto em = run(st, /*warm=*/false, req);
+    check(
+        "sampling: invalid direct worker value rejected before prefill",
+        em.threw && fake(st).prefill_calls == 0);
+  }
+}
+
 void test_new_full_prefill() {
   auto st = makeState();
   fake(st).steps = {{10, "a", false, false}, {0, "", true, true}};
@@ -446,6 +514,8 @@ void test_reset_named_only_clears_on_success() {
 
 int main() {
   printf("worker_loop.h harness:\n");
+  test_sampling_config_forwarded();
+  test_invalid_sampling_config_rejected();
   test_new_full_prefill();
   test_exact_prefix_warm_suffix();
   test_mismatch_full_reset();

examples/llm_server/cpp/worker_loop.h

@@ -62,9 +62,11 @@
 
 #include <algorithm>
 #include <chrono>
+#include <cmath>
 #include <cstdint>
 #include <iostream>
 #include <iterator>
+#include <limits>
 #include <memory>
 #include <stdexcept>
 #include <string>
@@ -115,6 +117,23 @@ inline void worker_handle_request(
   LLMSession& session = *st.session;
   int64_t max_new = req.value("max_new_tokens", static_cast<int64_t>(-1));
   const float temperature = req.value("temperature", 0.0f);
+  const double top_p_value = req.value("top_p", 1.0);
+  const int64_t top_k_value = req.value("top_k", static_cast<int64_t>(0));
+  const int64_t seed_value = req.value("seed", static_cast<int64_t>(0));
+  if (!std::isfinite(top_p_value) || top_p_value <= 0.0 || top_p_value > 1.0) {
+    throw std::runtime_error("top_p must be finite and in (0, 1]");
+  }
+  if (top_k_value < 0 || top_k_value > std::numeric_limits<int32_t>::max()) {
+    throw std::runtime_error("top_k must fit in a nonnegative int32");
+  }
+  // seed == 0 means "unset" (SamplingConfig::seed == 0 -> worker picks a random
+  // seed), so 0 is intentionally accepted here even though the HTTP layer
+  // treats 0 as the omitted sentinel and rejects an explicit seed=0. A JSON
+  // seed >= 2^63 won't fit int64_t and is rejected at parse time above; the
+  // HTTP layer caps the same range at 2^63 - 1 with a structured error.
+  if (seed_value < 0) {
+    throw std::runtime_error("seed must be nonnegative");
+  }
   // Stop strings (the request's `stop` sequences): terminate at the token
   // boundary where one appears so we don't generate to EOS/max_new past it. The
   // control plane also enforces these as a backstop.
@@ -207,6 +226,9 @@ inline void worker_handle_request(
 
   SamplingConfig sampling;
   sampling.temperature = temperature;
+  sampling.top_p = static_cast<float>(top_p_value);
+  sampling.top_k = static_cast<int32_t>(top_k_value);
+  sampling.seed = static_cast<uint64_t>(seed_value);
   const auto prefill_start = std::chrono::steady_clock::now();
   if (session.prefill_tokens(to_prefill, &sampling) !=
       ::executorch::runtime::Error::Ok) {

examples/llm_server/python/protocol.py

@@ -63,9 +63,6 @@ class ChatCompletionRequest(BaseModel):
     stop: Optional[Union[str, list[str]]] = None
     n: int = 1
     seed: Optional[int] = None
-    # Sampling knobs that change generation output. We don't plumb these, so they
-    # are modeled (not dropped) in order to be rejected with a clear error rather
-    # than silently ignored — see serving_chat's unsupported-parameter check.
     frequency_penalty: Optional[float] = None
     presence_penalty: Optional[float] = None
     top_k: Optional[int] = None

examples/llm_server/python/serving_chat.py

@@ -266,6 +266,9 @@ def _options(
         return GenerationOptions(
             max_new_tokens=req.resolved_max_tokens(),
             temperature=req.temperature if req.temperature is not None else 0.0,
+            top_p=req.top_p if req.top_p is not None else 1.0,
+            top_k=req.top_k if req.top_k is not None else 0,
+            seed=req.seed if req.seed is not None else 0,
             # Worker stop set, chosen per path in create() (see __init__ for the
             # two sets); the server re-applies it in _clean/_collect_until_stop.
             stop=stops,
@@ -347,6 +350,29 @@ def _reject_invalid_values(req: ChatCompletionRequest) -> None:
                 "invalid_request_error",
                 "invalid_value",
             )
+        if req.top_p is not None and (
+            not math.isfinite(req.top_p) or req.top_p <= 0.0 or req.top_p > 1.0
+        ):
+            raise APIError(
+                400,
+                f"top_p must be between 0 (exclusive) and 1 (got {req.top_p}).",
+                "invalid_request_error",
+                "invalid_value",
+            )
+        if req.top_k is not None and not (0 <= req.top_k <= 2**31 - 1):
+            raise APIError(
+                400,
+                f"top_k must be between 0 and {2**31 - 1} (got {req.top_k}).",
+                "invalid_request_error",
+                "invalid_value",
+            )
+        if req.seed is not None and not (0 < req.seed <= 2**63 - 1):
+            raise APIError(
+                400,
+                f"seed must be between 1 and {2**63 - 1} (got {req.seed}).",
+                "invalid_request_error",
+                "invalid_value",
+            )
         # max_tokens / max_completion_tokens, if given, must be positive integers
         # (OpenAI rejects 0 and negatives; our -1 sentinel means "unset/auto").
         for field_name in ("max_tokens", "max_completion_tokens"):
@@ -362,20 +388,17 @@ def _reject_invalid_values(req: ChatCompletionRequest) -> None:
     @staticmethod
     def _reject_unsupported_params(req: ChatCompletionRequest) -> None:
         """Reject params we don't honor rather than silently ignoring them (a
-        client relying on e.g. top_p/seed/logprobs would otherwise get wrong
-        behavior). Only the no-op/default value of each passes: top_p exactly
-        1.0; penalties 0; response_format type "text"; tool_choice none/auto/
-        unset; parallel_tool_calls true (false can't be guaranteed without
+        client relying on e.g. penalties/logprobs would otherwise get wrong
+        behavior). Only the no-op/default value of each passes: penalties 0;
+        response_format type "text"; tool_choice none/auto/unset;
+        parallel_tool_calls true (false can't be guaranteed without
         constraining); logprobs are not returned at all."""
         rf = req.response_format
         flags = [
             (req.n != 1, "n>1"),
-            (req.top_p is not None and req.top_p != 1.0, "top_p"),
-            (req.seed is not None, "seed"),
             (req.reasoning_effort is not None, "reasoning_effort"),
             (bool(req.frequency_penalty), "frequency_penalty"),
             (bool(req.presence_penalty), "presence_penalty"),
-            (req.top_k is not None, "top_k"),
             (bool(req.logit_bias), "logit_bias"),
             (
                 bool(rf) and rf.get("type", "text") != "text",
@@ -394,8 +417,8 @@ def _reject_unsupported_params(req: ChatCompletionRequest) -> None:
             raise APIError(
                 400,
                 f"Unsupported parameter(s): {', '.join(unsupported)}. This server honors "
-                "temperature, max_tokens/max_completion_tokens, stop, and tools for the "
-                "configured tool-call format.",
+                "temperature, top_p, top_k, seed, max_tokens/max_completion_tokens, "
+                "stop, and tools for the configured tool-call format.",
                 "invalid_request_error",
                 "unsupported_parameter",
             )

examples/llm_server/python/session_runtime.py

@@ -55,6 +55,9 @@ class GenerationOptions:
 
     max_new_tokens: int
     temperature: float = 0.0
+    top_p: float = 1.0
+    top_k: int = 0
+    seed: int = 0
     stop: list[str] = field(default_factory=list)
 
 
@@ -88,6 +91,9 @@ class GenStats:
 class _WorkerRequest:
     max_new_tokens: int
     temperature: float
+    top_p: float
+    top_k: int
+    seed: int
     stop: list[str]
     session_id: Optional[str]
     prompt_segments: Optional[list]
@@ -223,6 +229,9 @@ async def generate_stream(
         req = _WorkerRequest(
             max_new_tokens=options.max_new_tokens,
             temperature=options.temperature,
+            top_p=options.top_p,
+            top_k=options.top_k,
+            seed=options.seed,
             stop=list(options.stop),
             session_id=session_id,
             prompt_segments=prompt.segments,

examples/llm_server/python/tests/test_contract.py

@@ -96,7 +96,7 @@ def test_chat_streaming_protocol(make_client):
 
 
 def test_request_params_forwarded_to_generation(make_client):
-    # Contract behavior: the server must honor max_tokens/temperature.
+    # Contract behavior: the server must honor all supported sampling controls.
     client, fake = make_client()
     client.post(
         "/v1/chat/completions",
@@ -105,10 +105,16 @@ def test_request_params_forwarded_to_generation(make_client):
             "messages": [{"role": "user", "content": "hi"}],
             "max_tokens": 7,
             "temperature": 0.1,
+            "top_p": 0.8,
+            "top_k": 32,
+            "seed": 123,
         },
     )
     assert fake.captured_config.max_new_tokens == 7
     assert abs(fake.captured_config.temperature - 0.1) < 1e-6
+    assert abs(fake.captured_config.top_p - 0.8) < 1e-6
+    assert fake.captured_config.top_k == 32
+    assert fake.captured_config.seed == 123
 
 
 def test_special_tokens_forwarded_to_worker_as_stops(make_client):

examples/llm_server/python/tests/test_sampling_params.py

@@ -37,13 +37,9 @@ def test_n_greater_than_one_is_rejected(make_client):
 def test_unsupported_params_rejected(make_client):
     client, _ = make_client()
     for param in (
-        {"top_p": 0.5},
-        {"top_p": 2.0},  # > 1.0 is not a no-op either — only exactly 1.0/unset is
-        {"seed": 42},
         {"reasoning_effort": "high"},
         {"frequency_penalty": 1.0},
         {"presence_penalty": -0.5},
-        {"top_k": 40},
         {"logit_bias": {"123": 5.0}},
         {"response_format": {"type": "json_object"}},
         {"logprobs": True},
@@ -83,6 +79,23 @@ def test_temperature_range_rejected(make_client):
         assert r.json()["error"]["code"] == "invalid_value", param
 
 
+def test_sampling_parameter_ranges_rejected(make_client):
+    client, _ = make_client()
+    for param in (
+        {"top_p": 0.0},
+        {"top_p": -0.1},
+        {"top_p": 1.1},
+        {"top_k": -1},
+        {"top_k": 2**31},
+        {"seed": 0},
+        {"seed": -1},
+        {"seed": 2**63},
+    ):
+        r = client.post("/v1/chat/completions", json=_body(**param))
+        assert r.status_code == 400, param
+        assert r.json()["error"]["code"] == "invalid_value", param
+
+
 def test_noop_output_contract_fields_accepted(make_client):
     # The default/no-op forms must NOT be rejected (don't break OpenAI clients
     # that send them explicitly).
@@ -133,13 +146,19 @@ def test_unsupported_tool_choice_rejected(make_client):
 
 
 def test_supported_params_accepted(make_client):
-    # top_p=1.0 (no-op) and temperature/max_tokens must NOT be rejected; neither
+    # Sampling controls and temperature/max_tokens must not be rejected; neither
     # should tool_choice "auto" / "none".
     client, _ = make_client()
     for temperature in (0.0, 1.0, 2.0):
         r = client.post(
             "/v1/chat/completions",
-            json=_body(top_p=1.0, temperature=temperature, max_tokens=8),
+            json=_body(
+                top_p=0.9,
+                top_k=40,
+                seed=42,
+                temperature=temperature,
+                max_tokens=8,
+            ),
         )
         assert r.status_code == 200, temperature
     for choice in ("auto", "none"):