Add scheduler for batched requests

[metascroy]

This PR adds extension/llm/batching/, the scheduling seam for batched LLM serving. It is three headers. types.h holds the vocabulary shared by the runner, scheduler, and executor: an Input is one slice of work for one session, either a prompt chunk or a decode token, and a Task adds the identity used to order and cancel it. scheduler.h is the abstract interface (submit, has_work, get_work, cancel, clear), which is bookkeeping only: it never calls an executor or runs a callback, and it hands tasks back for the caller to complete. decode_first_scheduler.h is the first policy. It fills a batch with queued decodes first, at most one per session, then spends the leftover token budget on prefill, rotating one chunk per session so a long prompt cannot crowd out a decode. Everything is header-only and free of ExecuTorch runtime types, so the CMake target is an INTERFACE library that becomes a real one when the runner lands with a .cpp. Nothing consumes this yet; the runner is the next PR. Tested by 50 gtest cases in scheduler_test.cpp.

[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! (1 Unrelated Failure)

As of commit ae14128 with merge base 5c11944 ():

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

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

• pull / android / build-android (gh) (trunk failure)

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

[digantdesai]

not sure I follow this

[metascroy]

max_batch_size is the batch size forwarded to the model per step.

This scheduler admits enough work for max_decode_sequences_ and 2 prefill chunks per step. Decode work is always preferred to prefill work.

[digantdesai]

ok for now, but going forward I would like to separate scheduler mechanism and scheduler policy. I.e. Mechanism: how the scheduler interface looks like, Policy: how one can use those APIs to achieve a certain behavior, and we may impl more than one policy on the same mechanism. You can imagine on a phone we have a different policy vs. a beefy desktop.

[metascroy]

Absolutely! Focus on the public API (submit/has_work/get_work) for what the interface looks like (and even this I'm planning on modifying a bit as I'm sketching a model runner with concurrent sessions). We can promote to an interface, with this being just one (policy) implementation of that interface. A policy basically is 1-1 with the implementation of the get_work function, and the implementation of get_work here is just an example of a policy.

[metascroy]

To make this clearer, I split out scheduler.h into an interface, with decode_first_scheduler.h a specific implementation / policy.

[metascroy]

@claude review this code

[claude]

Claude finished @metascroy's task in 4m 53s —— View job

---

Code review: extension/llm/batching/ (post-refactor)

I re-read the split you made — scheduler.h (abstract mechanism), decode_first_scheduler.h (the decode-first policy), types.h, the CMake wiring, and the gtest suite. The mechanism/policy separation lands well and directly answers @digantdesai's note: Scheduler is now a pure interface, DecodeFirstScheduler is one policy behind it, and the static_asserts in the test lock that relationship in. A second policy can now sit on the same interface with no churn to callers. Nice.

My earlier findings are resolved by this refactor — worth recording since they no longer apply:

• The int32 truncation in n_tokens() is gone: sizes are std::size_t end to end (admissible_ compares input.size > max_prefill_chunk_size_, spent/n are size_t), so no signed truncation drives the budget.
• The retained-tokens-while-in-flight duplication is gone: get_work() now moves the Task out (std::move(*t)) and dispatch_ erases it from pending_, so nothing lingers.
• The complete()/fail() overload sharp edge is gone: settlement is now the caller's job; the scheduler is bookkeeping-only.

I traced the load-bearing invariants again and they still hold: the create() floor (max_prefill_chunk_size ≤ max_batch_tokens − max_decode_sequences) guarantees a full chunk fits on the first prefill pass after a saturated decode batch, so prefill is delayed but never permanently starved; rotation preserves "got-nothing outranks got-something"; and the rotation ⇔ map pairing stays duplicate-free across cancel/rejoin (the RejoiningAfterCancelDoesNotTakeTwoTurnsPerPass regression test is a good guard).

A few minor points, none blocking:

1. types.h uses bare size_t and omits <cstddef> (types.h:49-50, 73-76)
The other two headers use std::size_t and include <cstddef>; types.h uses unqualified size_t with only <cstdint>/<memory>/<optional>/<vector> included. It compiles today because <vector> transitively drags in size_t, but it's a portability footgun and inconsistent with its siblings. Add <cstddef> and qualify as std::size_t. Fix this →

2. Duplicated comment in clear() (scheduler.h:59-60)
// Drops every queued task and returns them all, for shutdown. is immediately followed by a stray // For shutdown. Fix this →

3. Move-then-pop ordering is inconsistent between the two passes (decode_first_scheduler.h:265-274 vs 328-332)
take_prefill_pass_ moves the task into taken before dq.pop_front(), but take_decodes_ pops from decode_queue_ before the taken.push_back(std::move(*t)). If that push_back ever throws (only bad_alloc realistically, and -fno-exceptions size builds are immune), the decode is gone from decode_queue_ yet still in pending_ — stranded: has_work() reports it but get_work() can never return it. Purely theoretical, but aligning take_decodes_ to the prefill path's "into-batch before pop" ordering would make it uniformly exception-safe. Minor. Fix this →

4. Non-copyable is implicit only (decode_first_scheduler.h:42) — DecodeFirstScheduler is correctly non-copyable/movable via its std::mutex, but an explicit = delete on the copy/move ops would document the intent. Optional.

The test suite remains the strongest part — the randomized never-lost/never-twice model, the fairness accounting, the cancelled-in-the-middle cases, and the two concurrency tests give real confidence. Overall this is a clean, well-reasoned seam; items 1–3 are the only ones I'd bother with, and all are small.
· add-scheduler