[executorch][cuda] Add a batched-sequence CUDA KV cache by Gasoonjia · Pull Request #24073 · pytorch/executorch

Gasoonjia · 2026-10-22T06:30:24Z

Stack from ghstack (oldest at bottom):

Adds CudaBatchedSequenceCache, registered as cache::kind::kBatchedSequence for the CUDA backend, the CUDA counterpart of MLX's batched sequence cache. Each sequence admitted with seq_new gets its own column of every layer's [rows, max_seqs + 1, H, D] storage (last column is scratch for padded static-width writes). Flat layers grow by whole rows; ring layers are fixed at min(capacity, window + max_write - 1) rows.

It implements the neutral BatchControl surface the batching executor already drives: declare_step / prepare_step / commit_step, seq_clone (device copy of the source prefix), seq_rm, rewind (bounded by the ring floor), pos, max_seqs. The per-sequence position plans reuse SequenceCache. prepare_step stages seq_cols and the seq_info segment table for triton::sdpa_varlen. DFlash owner blocks are validated as one run per owner.

CacheConfig gains max_seqs, the number of storage columns fixed at export. CudaKVPool exposes layer_data() for clone copies.

Differential Revision: D125987297

[ghstack-poisoned]

pytorch-bot · 2026-10-22T06:30:28Z

🔗 Helpful Links

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

📄 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! (10 Unrelated Failures)

As of commit b10c418 with merge base 4eb094b ():

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

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

MLX / test-mlx / test-mlx (gh) (trunk failure)
[ FAILED ] Implementations/MLXSequenceControlTest.DFlashRejectsCellOrPersistsOnlyContext/Sequence, where GetParam() = 4-byte object <01-00 00-00>
MLX / test-mlx-llm-offgraph (google/gemma-4-E2B-it, gemma4-e2b, gemma4) / test-mlx-llm-offgraph-gemma4-e2b (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
MLX / test-mlx-llm-offgraph (unsloth/gemma-3-1b-it, gemma3-1b, gemma) / test-mlx-llm-offgraph-gemma3-1b (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
MLX / test-mlx-llm-offgraph (unsloth/Llama-3.2-1B-Instruct, llama-1b, llama3) / test-mlx-llm-offgraph-llama-1b (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
MLX / test-mlx-parakeet / test-mlx-parakeet (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
MLX / test-mlx-stories110m / test-mlx-stories110m (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
MLX / test-mlx-supertonic / test-mlx-supertonic (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 8
MLX / test-mlx-voxtral / test-mlx-voxtral (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 1
MLX / test-mlx-voxtral-realtime / test-mlx-voxtral-realtime (gh) (trunk failure)
RuntimeError: Command bash /Users/runner/work/_temp/exec_script failed with exit code 134
Test CUDA Builds / test-muse-glimmer-cuda-e2e-kquant-17gb-dflash-image-batching / linux-job (gh) (trunk failure)
##[error]received SIGINT

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

github-actions · 2026-10-22T06:31:15Z

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.

[ghstack-poisoned]

…V columns Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * #24074 * #24073 * __->__ #24072 Adds `triton::sdpa_varlen`, the attention kernel the CUDA batched-sequence KV cache needs. Each resident sequence owns one column of a `[rows, cols, H, D]` store, and a device segment table `seq_info Int32[S, 4] = (q_start, q_len, col, kv_end)` says which packed query rows belong to which column and how long that sequence is. Lengths are read on device, so one captured CUDA graph serves any mix of sequence lengths, as the MLX per-sequence SDPA does per group. Covers causal, non-causal (DFlash draft blocks), sliding window and ring storage, GQA, rows no segment covers (zeroed), and split-K for short steps. A symbolic query length takes the generic path without host-side length branches, so exporting it with a dynamic step width does not pin that width. It is a new op in its own file; `triton::sdpa` and its dispatch are untouched. Differential Revision: D125987296 Pull Request resolved: #24072

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * __->__ #24074 * #24073 * #24072 Teaches `LowerOffGraphKVPass` a third manifest layout, `batched-sequence`, alongside `sequence` (solo) and `cell`. Each `kvcache.update_and_attend` becomes an in-graph write to slot `(position % rows) * columns + seq_cols[i]` followed by `triton::sdpa_varlen` over the column store, with `seq_cols` / `seq_info` as runtime-owned step buffers. The manifest carries `max_seqs`, `max_segments` (= 2 * max_seqs) and `causal` (false for DFlash drafts). `CudaBackend` treats it as a packed layout like `cell` (step-width check from the packed position source; no joint-graph constant folding of runtime buffers). Differential Revision: D125987298 Pull Request resolved: #24074

…ence cache Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * __->__ #24075 * #24074 * #24073 * #24072 Adds `--kv-layout {cell,sequence}` and `--max-kv-seqs` to the Muse Glimmer CUDA batching exports (solo and DFlash). `sequence` writes a `batched-sequence` manifest for target and draft and publishes `get_offgraph_kv_max_seqs` (and `get_draft_offgraph_kv_max_seqs` / `get_draft_offgraph_kv_seq_capacity` for the draft). `CudaExecutor` and the CUDA DFlash executor pick `kBatchedSequence` when the artifact declares a sequence count, `kBatchedCell` otherwise, so the served feature set matches MLX's batched-sequence path. Cell artifacts are unchanged. The layout is chosen at export; switching it at runtime on one .pte is out of scope. Differential Revision: D125987299 Pull Request resolved: #24075

…the sampler rows Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * __->__ #24107 * #24075 * #24074 * #24073 * #24072 `forward_others` and `forward_from_embeddings` take up to `max_step_tokens` (4096) tokens and select which rows' logits to return. The selector's dynamic bound was also `max_step_tokens`, so memory planning sized each method's FP32 logits output at `4096 x 202048 x 4 B = 3157 MiB`, allocated at load whether or not a request ever uses it. A step never selects more rows than the device samplers take (`MAX_SAMPLE_ROWS = 64`: one row per prefilling session, or `max_owners x block` for DFlash verify), and the executor already rejects a step with more. Bound the selector at `min(max_step_tokens, MAX_SAMPLE_ROWS)`. Each planned logits output drops to 49 MiB; token width and the DFlash taps keep their `max_step_tokens` bound. Differential Revision: D127114350 Pull Request resolved: #24107

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * __->__ #24076 * #24107 * #24075 * #24074 * #24073 * #24072 A ring layer cannot wrap before it holds all its slots, so until then it grows with the flat layers instead of allocating window + max_write - 1 rows per sequence up front. Growable layers now cap at their own declared rows. Differential Revision: D126851225 Pull Request resolved: #24076

@Gasoonjia

This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #24073 by @Gasoonjia ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/91/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/91/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/90/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/91/orig Differential Revision: [D125987297](https://our.internmc.facebook.com/intern/diff/D125987297/) @diff-train-skip-merge Co-authored-by: gasoonjia <gasoonjia@icloud.com>

Update

7c52884

[ghstack-poisoned]

Gasoonjia requested review from larryliu0820 and mergennachin as code owners October 22, 2026 06:30

Gasoonjia temporarily deployed to cadence October 22, 2026 06:30 — with GitHub Actions Inactive

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 Oct 22, 2026

meta-codesync Bot added the meta-exported label Oct 22, 2026

Gasoonjia temporarily deployed to cadence October 22, 2026 06:48 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence October 22, 2026 06:49 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence October 22, 2026 06:52 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to upload-benchmark-results October 22, 2026 21:58 — with GitHub Actions Inactive

Update

280d275

[ghstack-poisoned]

Gasoonjia mentioned this pull request Oct 23, 2026

[executorch][muse-glimmer] Bound the CUDA prefill logits selector by the sampler rows #24107

Merged

Gasoonjia temporarily deployed to cadence October 23, 2026 08:51 — with GitHub Actions Inactive

Gasoonjia deployed to upload-benchmark-results November 6, 2026 01:41 — with GitHub Actions Active

Update

c495e3c

[ghstack-poisoned]

Gasoonjia temporarily deployed to cadence November 9, 2026 08:12 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence November 9, 2026 08:26 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence November 9, 2026 08:28 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence November 9, 2026 08:30 — with GitHub Actions Inactive

Gasoonjia deployed to upload-benchmark-results November 9, 2026 09:48 — with GitHub Actions Active

metascroy approved these changes Nov 9, 2026

View reviewed changes

Update

b10c418

[ghstack-poisoned]

Gasoonjia temporarily deployed to cadence November 10, 2026 05:06 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence November 10, 2026 05:23 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence November 10, 2026 05:26 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence November 10, 2026 05:27 — with GitHub Actions Inactive

Gasoonjia deployed to upload-benchmark-results November 10, 2026 11:05 — with GitHub Actions Active

meta-codesync Bot merged commit e2d24a7 into gh/Gasoonjia/91/base Nov 11, 2026
321 of 332 checks passed

meta-codesync Bot deleted the gh/Gasoonjia/91/head branch November 11, 2026 02:50

meta-codesync Bot deployed to cherry-pick-bot November 11, 2026 02:50 Active

pytorchbot mentioned this pull request Nov 11, 2026

[executorch][cuda] Add a batched-sequence CUDA KV cache #24880

Merged

Navigation Menu

Repository navigation

[executorch][cuda] Add a batched-sequence CUDA KV cache - #24073

[executorch][cuda] Add a batched-sequence CUDA KV cache#24073
meta-codesync[bot] merged 5 commits into
gh/Gasoonjia/91/basefrom
gh/Gasoonjia/91/head

Gasoonjia commented Oct 22, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Oct 22, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 22, 2026

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

Gasoonjia commented Oct 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Oct 22, 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/24073

✅ You can merge normally! (10 Unrelated Failures)

Uh oh!

github-actions Bot commented Oct 22, 2026

This PR needs a release notes: label

Uh oh!

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Gasoonjia commented Oct 22, 2026 •

edited

Loading

pytorch-bot Bot commented Oct 22, 2026 •

edited

Loading

This PR needs a `release notes:` label