[executorch][muse-glimmer] Bound the CUDA prefill logits selector by the sampler rows by Gasoonjia · Pull Request #24107 · pytorch/executorch

Gasoonjia · 2026-10-23T08:50:34Z

Stack from ghstack (oldest at bottom):

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

[ghstack-poisoned]

pytorch-bot · 2026-10-23T08:50:39Z

🔗 Helpful Links

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

📄 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, 10 Unrelated Failures

As of commit a0e8579 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-23T08:51:18Z

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.

…ckends Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * #24074 * #24073 * #24072 * #24066 * #24065 * #24064 * #24063 * __->__ #24062 Move the MLX executor's EngineEmbeddings (vision encoder + embed_text + MuseGlimmerMaterializer) into runtime/model_embeddings so the CUDA batching executors can materialize prepared image inputs the same way. MLX behavior is unchanged; it now holds a MuseGlimmerModelEmbeddings. Differential Revision: D125987290 Pull Request resolved: #24062

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * #24074 * #24073 * #24072 * #24066 * #24065 * #24064 * __->__ #24063 * #24062 Steps that carry a prepared (image) input run forward_from_embeddings with the materialized rows; token-only steps keep the static forward_{N} methods and their CUDA graphs. SessionTable packs prepared rows as placeholders so the cache still plans their positions. Differential Revision: D125987291 Pull Request resolved: #24063

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * #24074 * #24073 * #24072 * #24066 * #24065 * __->__ #24064 * #24063 * #24062 export_cuda_batching optionally adds vision_encoder, embed_text and a dynamic forward_from_embeddings over the same off-graph KV contract, and writes pos_embed.bin next to the program. Differential Revision: D125987292 Pull Request resolved: #24064

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * #24074 * #24073 * #24072 * #24066 * __->__ #24065 * #24064 * #24063 * #24062 The CUDA backend factory builds the image preparation spec from the shared embeddings engine, and serve_batching passes --pos-embed-path through, so image requests are served on regular CUDA batching. Differential Revision: D125987293 Pull Request resolved: #24065

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom): * #24076 * #24107 * #24075 * #24074 * #24073 * #24072 * __->__ #24066 * #24065 * #24064 * #24063 * #24062 The CUDA DFlash executor accepts prepared image inputs on admission; image steps run the target from embeddings, then hidden taps and drafting proceed as for tokens. Differential Revision: D125987295 Pull Request resolved: #24066

[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

…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

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

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

Update

dedcd5f

[ghstack-poisoned]

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

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

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

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

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

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

Update

5887d81

[ghstack-poisoned]

Gasoonjia temporarily deployed to cadence November 5, 2026 22:02 — with GitHub Actions Inactive

metascroy approved these changes Nov 5, 2026

View reviewed changes

Gasoonjia temporarily deployed to cadence November 5, 2026 22:16 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence November 5, 2026 22:20 — with GitHub Actions Inactive

Gasoonjia temporarily deployed to cadence November 5, 2026 22:23 — with GitHub Actions Inactive

Update

2cf8b24

[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:27 — 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:31 — with GitHub Actions Inactive

Update

a0e8579

[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

meta-codesync Bot merged commit 2d090bb into gh/Gasoonjia/95/base Nov 11, 2026
266 of 277 checks passed

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

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

pytorchbot mentioned this pull request Nov 11, 2026

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

Merged

Navigation Menu

Repository navigation

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

[executorch][muse-glimmer] Bound the CUDA prefill logits selector by the sampler rows#24107
meta-codesync[bot] merged 4 commits into
gh/Gasoonjia/95/basefrom
gh/Gasoonjia/95/head

Gasoonjia commented Oct 23, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Oct 23, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 23, 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 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot Bot commented Oct 23, 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/24107

❌ 1 New Failure, 10 Unrelated Failures

Uh oh!

github-actions Bot commented Oct 23, 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 23, 2026 •

edited

Loading

pytorch-bot Bot commented Oct 23, 2026 •

edited

Loading

This PR needs a `release notes:` label