Add scheduler for batched requests

[metascroy]

Adds extension/llm/scheduler, a header-only step scheduler for batched LLM serving. A Request is one step for one sequence, either a single decode token or one prefill chunk, which callers submit() and await via std::future, with get_work() assembling the next Batch. Decodes are served first in arrival order up to max_decode_sequences, then the remainder of max_batch_size goes to prefill, which is organized as a FIFO per session plus a rotation across sessions so a long prompt advances one chunk per pass instead of monopolizing the batch. The scheduler reads only request_id, session_id, and tokens.size(). Everything else lives in RequestParams/ResponsePayload and is carried through untouched, so payload fields can be added without touching scheduling logic. SchedulerParams is immutable and self-validating with max_batch_size derived rather than stored, making inconsistent limits unrepresentable. The payload already accommodates upcoming dflash speculative decoding: OutputRows::All requests a prediction at every drafted position, an absent sampling returns the raw distribution instead of tokens, and accept/reject stays in the caller where the draft output lives. Ships with 53 gtest cases covering scheduling order, fairness, cancellation, completion, and a randomized invariant check that models queued() exactly. All pass under -Wall -Wextra -Werror and ThreadSanitizer, and the suite is wired into CMake as an INTERFACE target beside extension_llm_cache.

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

✅ You can merge normally! (6 Unrelated Failures)

As of commit b7131d7 with merge base 89e6f90 ():

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• pull / unittest / linux / linux-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest / macos / macos-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest / windows / windows-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / linux / linux-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / macos / macos-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions
• pull / unittest-editable / windows / windows-job (gh) (trunk failure)
  .ci/scripts/tests/test_filter_cuda_matrix.py::TestPublishedSets::test_published_python_versions

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

[metascroy]

@claude review this code

[github-actions]

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 finished @metascroy's task in 2m 16s —— View job

---

Code review: extension/llm/scheduler

I read scheduler.h, the 53 gtest cases, and the CMake wiring. This is high-quality, well-reasoned code — the invariants are documented where they're load-bearing, the exception-safety ordering (copy-into-batch before pop) is deliberate and correct, and the test suite is genuinely thorough (fairness, cancellation-in-the-middle, queued_ underflow, stale-rotation entries, the randomized invariant model). I traced the trickiest parts and they hold up:

• Rotation fairness — deferred (didn't fit) are re-inserted front-in-original-order, served go to the back, so "got-nothing outranks got-something" is preserved. Verified across the mixed deferred/served/drained cases.
• No permanent prefill starvation — max_batch_size = 2·chunk + decodes guarantees a full chunk fits on the first prefill pass after decodes are taken. This is the "load-bearing" comment and it checks out.
• rotation ⇔ map invariant — no duplicate rotation entries, and the submit catch block leaving a stale rotation entry is genuinely harmless (get_work() drops an entry with no session).
• complete/fail double-settle — both find+erase under the same lock before touching the promise, so exactly one side wins. Correct.

A few minor points, none blocking:

1. Unvalidated int32 truncation in Request::n_tokens() (scheduler.h:122-124)
static_cast<std::int32_t>(tokens.size()) can go negative for a >2³¹-token vector. A decode is guarded by size()==1, but a prefill with a truncated-negative n would slip past the n > max_prefill_chunk_size() check (negative < positive) and then corrupt the budget in take_prefill_pass_. Purely theoretical at these sizes, but it's the one place an untrusted size() becomes a signed quantity that drives arithmetic. A guard in submit() (reject tokens.size() > max_prefill_chunk_size() on the size_t before casting) would close it. Fix this →

2. Retained request keeps a dead tokens copy while in flight (scheduler.h:415, 453)
get_work() copies the Request (tokens included) into the Batch, but the PendingRequest in pending_requests_ also retains the full request until complete()/fail(). After dispatch, settle_ only ever reads p->request.session_id — the retained tokens vector (up to max_prefill_chunk_size per in-flight request) is dead weight. Clearing it->second->request.tokens after the successful push_back would drop that duplication without disturbing the intentional copy-before-pop ordering. Minor.

3. complete() overload sharp edge (scheduler.h:299, 306)
The two overloads (std::vector<Token> vs LogitsPtr) are unambiguous for typed arguments, but complete(id, {}) or complete(id, nullptr) could be ambiguous / silently settle with a null LogitsPtr. Not a bug given how it's meant to be called; worth a one-line note that a null logits pointer isn't validated.

4. Non-copyable is implicit only (scheduler.h:216)
Scheduler is correctly non-copyable/movable because it holds a std::mutex/std::atomic, but that's implicit. An explicit = delete on the copy/move ops would document the intent for future readers. Optional.

Nothing here changes the design; items 1 and 2 are the only ones I'd consider acting on, and both are minor. Nice work.
· add-scheduler