[executorch][muse] Serve Muse Glimmer CUDA batching on a batched-sequence cache by Gasoonjia · Pull Request #24075 · pytorch/executorch

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

Stack from ghstack (oldest at bottom):

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

[ghstack-poisoned]

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

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

❌ 1 New Failure, 1 Pending, 10 Unrelated Failures

As of commit 7278b82 with merge base 4eb094b ():

NEW FAILURE - The following job has failed:

pull / unittest / macos / macos-job (gh)
examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity

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:31Z

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

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

…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

…ence cache (#24882) This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #24075 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/93/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/93/head Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/92/orig Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/93/orig Differential Revision: [D125987299](https://our.internmc.facebook.com/intern/diff/D125987299/) @diff-train-skip-merge Co-authored-by: gasoonjia <gasoonjia@icloud.com>

shoumikhin · 2026-11-11T09:32:45Z

heads up, the new test_cuda_sequence_export_inputs.py broke unittest / macos on main (the editable and release macos rows too). it imports the cuda backend, which pulls in triton, and the macos runners have no triton:

FAILED examples/models/muse-glimmer/tests/test_cuda_sequence_export_inputs.py::SequenceExportInputsTest::test_small_logits_example_preserves_token_and_embedding_capacity - ModuleNotFoundError: No module named 'triton'

same failure shows on this pr's own macos run. probably easiest to skip the test when triton isn't installed, or add it to the muse-glimmer --ignore list in pytest.ini next to the other test_cuda_* files. it's blocking viable/strict, so a quick fix would help

Update

28effa3

[ghstack-poisoned]

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

This was referenced Oct 22, 2026

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

Merged

[executorch][cuda] Lower off-graph KV to the batched-sequence layout #24074

Merged

[executorch][cuda] Grow sequence-layout ring layers on demand #24076

Merged

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

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

Gasoonjia temporarily deployed to cadence October 22, 2026 06:50 — 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 22:54 — with GitHub Actions Inactive

Update

01a1fde

[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

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

Gasoonjia temporarily deployed to cadence October 23, 2026 09:12 — with GitHub Actions Inactive

Update

9fb23e4

[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:29 — 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:41 — with GitHub Actions Active

metascroy approved these changes Nov 9, 2026

View reviewed changes

Update

7278b82

[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:24 — 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 12:57 — with GitHub Actions Active

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

meta-codesync Bot deleted the gh/Gasoonjia/93/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][muse] Serve Muse Glimmer CUDA batching on a batched-sequence cache #24882

Merged

Navigation Menu

Repository navigation

[executorch][muse] Serve Muse Glimmer CUDA batching on a batched-sequence cache - #24075

[executorch][muse] Serve Muse Glimmer CUDA batching on a batched-sequence cache#24075
meta-codesync[bot] merged 5 commits into
gh/Gasoonjia/93/basefrom
gh/Gasoonjia/93/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!

shoumikhin commented Nov 11, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 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/24075

❌ 1 New Failure, 1 Pending, 10 Unrelated Failures

Uh oh!

github-actions Bot commented Oct 22, 2026

This PR needs a release notes: label

Uh oh!

Uh oh!

shoumikhin commented Nov 11, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Gasoonjia commented Oct 22, 2026 •

edited

Loading

pytorch-bot Bot commented Oct 22, 2026 •

edited

Loading

This PR needs a `release notes:` label