Repository navigation
[executorch][cuda] Grow sequence-layout ring layers on demand - #24076
Merged
Merged
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24076
Note: Links to docs will display an error until the docs builds have been completed. ❌ 2 New Failures, 1 Cancelled JobAs of commit 9003ff7 with merge base 4eb094b ( NEW FAILURES - The following jobs have failed:
CANCELLED JOB - The following job was cancelled. Please retry:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
This was referenced Oct 21, 2026
meta-cla
Bot
added
the
CLA Signed
This PR needs a
|
Gasoonjia
temporarily deployed
to
upload-benchmark-results
October 22, 2026 23:26 — with 
GitHub Actions
Inactive
GitHub Actions
Inactive
Gasoonjia
added a commit
that referenced
this pull request
Oct 23, 2026
Pull Request resolved: #24076 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. ghstack-source-id: 455353645 @exported-using-ghexport Differential Revision: [D126851225](https://our.internmc.facebook.com/intern/diff/D126851225/)
Gasoonjia
added a commit
that referenced
this pull request
Nov 9, 2026
Pull Request resolved: #24076 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. ghstack-source-id: 468443991 @exported-using-ghexport Differential Revision: [D126851225](https://our.internmc.facebook.com/intern/diff/D126851225/)
Gasoonjia
added a commit
that referenced
this pull request
Nov 10, 2026
Pull Request resolved: #24076 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. ghstack-source-id: 469359765 @exported-using-ghexport Differential Revision: [D126851225](https://our.internmc.facebook.com/intern/diff/D126851225/)
meta-codesync Bot
pushed a commit
that referenced
this pull request
Nov 11, 2026
…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
meta-codesync Bot
pushed a commit
that referenced
this pull request
Nov 11, 2026
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
meta-codesync Bot
pushed a commit
that referenced
this pull request
Nov 11, 2026
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * #24074 * __->__ #24073 * #24072 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 Pull Request resolved: #24073
meta-codesync Bot
pushed a commit
that referenced
this pull request
Nov 11, 2026
…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
meta-codesync Bot
pushed a commit
that referenced
this pull request
Nov 11, 2026
…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
meta-codesync
Bot
merged commit Nov 11, 2026
fd03531
into
gh/Gasoonjia/94/base
278 of 282 checks passed
Gasoonjia
added a commit
that referenced
this pull request
Nov 11, 2026
This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #24076 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/94/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/94/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/95/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/94/orig Differential Revision: [D126851225](https://our.internmc.facebook.com/intern/diff/D126851225/) @diff-train-skip-merge Co-authored-by: gasoonjia <gasoonjia@icloud.com>
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Stack from ghstack (oldest at bottom):
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