Add a cell-cache executor for batched generation by kiymetakdemir · Pull Request #22234 · pytorch/executorch

kiymetakdemir · 2026-08-27T19:28:14Z

Implements the Executor seam over the cell KV cache. Backend-neutral: it reaches the cache through CacheBuilderRegistry and BatchControl, with the backend id, cache kind and cache-key option as configuration, so any backend registering a cell cache is served. A session is one cache sequence, and a batch is one forward carrying every input's tokens end to end on a single axis, with the cache's mask keeping the sequences apart. SessionIds are monotonic and never reissued, mapped onto cache sequence ids that do recycle; the cache frees an id once a sequence's last cell goes, which the seam forbids for a SessionId.

Read build_step first; it holds the reasoning. It flattens a batch into tokens, positions and sequence ids in one pass so entry i of each names the same token, which is how the cache pairs them when placing cells and building the mask. Sequence lengths come from the cache and do not move as inputs are laid down, so a per-sequence cursor carries the batch's own
writes.

The method loads on the first execute() rather than in create(), because a delegate may bind per-thread state as it initializes, which happens during that load, and construction is the only entry point that does not run on the engine thread.

pytorch-bot · 2026-08-27T19:28:19Z

🔗 Helpful Links

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

📄 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 cb15008 with merge base 5c11944 ():

FLAKY - The following job failed but was likely due to flakiness present on trunk:

pull / test-qnn-testsuite-linux / test-backend-linux (qnn, models) / linux-job (gh) (detected as infra flaky with no log or failing log classifier)

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

metascroy · 2026-08-27T23:46:05Z

+        // The sampler reduces in place, so copy rather than overwrite the
+        // model's output.
+        std::vector<float> scores(begin, begin + vocab);
+        Sampler sampler(


Don't construct a new sampler on each execute.

We should have session_id -> SessionInfo map

And the SessionInfo can have a sampler field on it that gets set during set_sampling_params or something like that. We also don't need to do all this seed mixing.

We can just continue using the sampler object.

Makes sense, thank you! Sampler needs vocab_size at construction and set_sampling doesn't have it. The export script can save that to metadata. I guess the other option is building the sampler lazily on the first draw.

We don't need to build sampler lazily.

The exeuctor/model should know its vocab size, and can save in private field after construction that sampler construction reads from

metascroy · 2026-08-27T23:52:05Z

+
+// Spread one seed over positions so a session's neighbouring tokens draw
+// unrelated streams.
+std::uint64_t mix(std::uint64_t seed, std::int64_t position) {


Do we need mix? See comment below on sampler

Removing it.

metascroy · 2026-08-28T00:01:46Z

+  out.outputs.resize(batch.inputs.size());
+
+  // Ahead of the step: a refused load must claim no cell.
+  if (!ensure_method_loaded()) {


Why are we loadign in execute?

It was a work around for the stream issue I got. From what I understand, the caller constructs the executor and passes it to Runner by reference, so the executor exists before the Runner, and the Runner creates the engine thread. That means create() runs on the caller's thread while every execute() runs on the engine thread.
Loading in create() registers it on the caller's thread, so the first execute() fails with "There is no Stream(gpu, N) in current thread".
But again thinking about it, a start() on Executor, called once on the engine thread before any batch should solve the issue.

I thought your other PR where you do cache building lazily is what fixed stream issue?

They are actually two different stream issues; one is the model's own stream created on the wrong thread, the other is the KV pool using the default stream instead of the one it was handed.

What is a minimal repro for the second model stream on wrong thread?

That seems like a bigger UI issue that I don't want to paper over with a short term fix.

metascroy · 2026-08-28T00:04:25Z

+      make_tensor_ptr({1, static_cast<int>(step->tokens.size())}, step->tokens);
+  auto positions = make_tensor_ptr(
+      {static_cast<int>(step->positions.size())}, step->positions);
+  const auto result = module_->execute(config_.method, {tokens, positions});


What communicates logit positions to the model?

metascroy · 2026-08-28T00:06:04Z

+  const auto logits = result->at(0).toTensor();
+
+  for (std::size_t i = 0; i < batch.inputs.size(); ++i) {
+    const int row = step->logit_indices[i];


How are we communicating logit indices to the model in lm_head project?

We don't, we talked about changing the model forward but didn't do it yet, should I do it as a follow up pr?

We can do as follow-up PR, but I was very surprised it wasn't already being done at export stage

I guess I just focused on single sequence that time

But even for single sequence, I expected last token selection, but it seems that isn't being done

metascroy · 2026-08-28T00:09:32Z

+    // limits the table cannot honor and open_session() holds the count --
+    // exhaustion is kept unreachable rather than handled.
+    int max_sessions = 0;
+    int max_session_tokens = 0;


Is max_session_tokens tied to some limit in cacheconfig?

We are checking if max_sessions * max_session_tokens > cacheconfig.capacity and reject if it exceeds capacity.

metascroy · 2026-08-28T00:11:32Z

+      make_tensor_ptr({1, static_cast<int>(step->tokens.size())}, step->tokens);
+  auto positions = make_tensor_ptr(
+      {static_cast<int>(step->positions.size())}, step->positions);
+  const auto result = module_->execute(config_.method, {tokens, positions});


What is the shape of result?

Is it logits per token, or only final logits?

It is logits per token

metascroy · 2026-08-28T00:12:37Z

+      ET_CHECK_MSG(false, "CellExecutor: unsupported logits dtype");
+    }
+  } ctx;
+  ET_SWITCH_THREE_TYPES(


Can we not reuse sample_from_logits?

We can but sample_from_logits assumes single-sequence decoding. I can either add a row parameter defaulting to last, or pass a one-row tensor view so the offset branch is skipped.

Is it difficult to call sample_from_logits per sequence in batch?

No, I can do that

metascroy · 2026-08-28T00:19:17Z

+  // (backend_id, cache_kind).
+  static std::unique_ptr<CellExecutor> create(
+      std::unique_ptr<Module> module,
+      Config config);


Just have params in create directly? Do we need another Config struct?

Config has a CacheConfig inside of it, but CacheConfig is fixed by Module and the model archetecture. Can we just construct from Module? Do we have shareable helpers for that?

CacheConfig has parameters the module doesn't carry: capacity, kv_storage_dtype, initial_capacity; so create can't build it from the module alone, the application supplies those from its own flags.

But the executor just reads one field of it (capacity, for the max_sessions × max_session_tokens check) and forwards the rest to the cache builder.

Would it make more sense for the application to build the cache and pass it in? The executor would take capacity from the controller, and what create needs would be max_sessions, max_session_tokens, backend_id, and method.

I think there are two parts to cache config. There are parts that are specific to the model and not actually configurable (these are read from pte metadata), and then there are parts that are configurable.

Maybe CacheConfig is a class with methods on it and helper that builds it from module?

auto config = CacheConfig::create(model) or something. config.set_capacity(max_sessions × max_session_tokens)

What are your thoughts?

Makes sense to me, especially passing max_sessions, max_session_tokens and capacity separately and then comparing them looked redundant

metascroy · 2026-08-28T22:41:30Z

+  // once before any other call, on that thread.
+  //
+  // false = the executor cannot run.
+  virtual bool start() {


I'm still not understanding this threading issue you're hitting, and why you need this. What is the exact issue?

metascroy · 2026-08-28T22:41:58Z

+      std::string method = "forward");
+
+  // Loads the method, naming the cache to the backend. The delegate binds
+  // per-thread state as it initializes, which happens during that load.


Code pointer for this?

Module::load_method
→ Method::init runtime/executor/method.cpp:883
→ BackendDelegate::Init runtime/executor/method.cpp:974
→ backend->init(...) runtime/executor/method.cpp:123
→ new (handle) MLXHandle() MLXBackend.cpp:230
→ new_stream(Device::gpu) MLXBackend.cpp:194

metascroy · 2026-08-28T22:47:43Z

+}
+
+std::optional<Token> CellExecutor::sample_row(
+    ::executorch::aten::Tensor& logits,


Can we use the non-aten tensor?

metascroy · 2026-08-28T22:48:23Z

+    int max_sessions,
+    int max_session_tokens,
+    int kv_dtype,
+    std::string backend_id,


Doesn't module know the backend?

I now read it from the program and dropped the parameter, rejecting a method spanning more than one backend.

metascroy · 2026-08-28T22:52:07Z

+  if (!cfg.ok()) {
+    return nullptr;
+  }
+  cfg->capacity = max_sessions * max_session_tokens;


How does capacity differ from initial_capacity?

If capacity dynamically grows, why do we need to set it to max_sessions * max_session_tokens?

capacity is the ceiling so it doesn't grow, initial_capacity is the starting allocation. I set it to the product because that's the most the executor can ever need.

metascroy · 2026-08-28T22:54:24Z

+      logits_sizes[logits_sizes.size() - 1]));
+}
+
+bool CellExecutor::start() {


this looks like it should be called load?

Removed this after the mlx backend fix.

metascroy · 2026-08-28T22:58:33Z

+  // before the step, so the batch's own writes live here.
+  std::unordered_map<std::int32_t, int> cursor;
+
+  for (const Input& input : batch.inputs) {


Does model have a max accepted size via prefill chunk size or something?

Are we overflowing that limit?

…om it

Add a cell-cache executor for batched generation

3e31cf1

kiymetakdemir requested review from SS-JIA, digantdesai and metascroy August 27, 2026 19:28

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 27, 2026

metascroy reviewed Aug 27, 2026

View reviewed changes

metascroy reviewed Aug 28, 2026

View reviewed changes

Address review: build the cache from the program, add a start hook

ef6400a

metascroy reviewed Aug 28, 2026

View reviewed changes

Address review: build the cache from the program, read the backend fr…

cb15008

…om it

Uh oh!

Conversation

kiymetakdemir commented Aug 27, 2026

Uh oh!

pytorch-bot Bot commented Aug 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

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

✅ You can merge normally! (1 Unrelated Failure)

Uh oh!

Choose a reason for hiding this comment

Uh oh!

kiymetakdemir Aug 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pytorch-bot Bot commented Aug 27, 2026 •

edited

Loading

kiymetakdemir Aug 28, 2026 •

edited

Loading