Skip to content

[executorch][cuda] Test CudaExecutor on every step width up to 65 - #23819

Merged
meta-codesync[bot] merged 6 commits into
gh/Gasoonjia/59/basefrom
gh/Gasoonjia/59/head
Oct 20, 2026
Merged

[executorch][cuda] Test CudaExecutor on every step width up to 65#23819
meta-codesync[bot] merged 6 commits into
gh/Gasoonjia/59/basefrom
gh/Gasoonjia/59/head

Conversation

@Gasoonjia

@Gasoonjia Gasoonjia commented Oct 14, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

CudaExecutor runs a step of T tokens through the narrowest static
forward_{N} that holds it, padding T up to N (for power-of-two widths, the
next power of two), and sends the padding rows' K/V to the pool's scratch row
and their logits rows nowhere; wider steps run forward_others unpadded, and
steps wider than it slice. The GPU tests drove only the toy's forward_{1,2,4,8}
through the runner, so the widths Muse Glimmer's forward_{16,32,64} pad to, and
the step just past the widest static method, were never run.

  • export_toy_decoder.py adds mg_widths/: Muse Glimmer's method set,
    forward_{1,2,4,...,64} and forward_others from 64 tokens selecting at least
    five rows, over eight prompts of 1..50 tokens with 12-token continuations.
  • EveryStepWidthUpTo65MatchesEagerGreedy: every width 1..65 (and 96, 97) six
    times, shuffled. Each step is a random mix of the sequences (each at most
    once) with random chunk lengths summing to the width, some chunks asking for
    their next token and some not; sequences are teacher-forced along their
    prompt and eager continuation. Each step must run the expected method(s)
    (forward_{next pow2}; forward_others for 65..96; forward_others +
    forward_1 for 97), return exactly the outputs asked for, and every output
    must be the eager greedy token at that position. Runs with CUDA graphs on and
    off.

Differential Revision: D125110873

@Gasoonjia
Update
43ddd32 
[ghstack-poisoned]
@pytorch-bot

pytorch-bot Bot commented Oct 14, 2026 •
edited
Loading

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 4 New Failures, 3 Cancelled Jobs, 1 Unclassified Failure

As of commit 0ff7b05 with merge base 897d28f (image):

NEW FAILURES - The following jobs have failed:

UNCLASSIFIED FAILURE - DrCI could not classify the following job because the workflow did not run on the merge base. The failure may be pre-existing on trunk or introduced by this PR:

CANCELLED JOBS - The following jobs were cancelled. Please retry:

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

@meta-cla 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 14, 2026
This was referenced Oct 14, 2026
Merged
Closed
Merged
Merged
Merged
Merged
Merged
Merged
@github-actions

github-actions Bot commented Oct 14, 2026

Copy link
Copy Markdown

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.

@Gasoonjia
Gasoonjia deployed to upload-benchmark-results October 14, 2026 17:32 — with GitHub Actions Active
@Gasoonjia
Update
af15189 
[ghstack-poisoned]
This was referenced Oct 18, 2026
Open
Open
Open
Open
Open
Open
Open
Open
Open
Open
@Gasoonjia
Update
0d927ed 
[ghstack-poisoned]
@Gasoonjia
Gasoonjia deployed to upload-benchmark-results October 19, 2026 09:28 — with GitHub Actions Active
@Gasoonjia
Update
e8dc8a0 
[ghstack-poisoned]
Gasoonjia added a commit that referenced this pull request Oct 19, 2026
@Gasoonjia
[executorch][cuda] Test CudaExecutor on every step width up to 65
d0897af 
Pull Request resolved: #23819

CudaExecutor chooses the narrowest static forward_{N} that holds a step, pads unused rows into the KV scratch cell without returning their logits, and uses forward_others with slicing for wider steps. Extend the toy artifact to the Muse-sized method set forward_{1,2,4,...,64} plus forward_others, with eight independently referenced prompts and continuations.

EveryStepWidthUpTo65MatchesEagerGreedy covers every width 1..65 plus 96 and 97, shuffled repeatedly, with CUDA graphs both on and off. Mix sequence owners and chunk lengths, teacher-force eager reference histories, and vary which chunks request output. Assert the selected method sequence, exact output ownership/count and eager greedy token at each requested position. This is correctness coverage of static padding, dynamic widths and slicing, without serving performance reporting.
ghstack-source-id: 451833966
@exported-using-ghexport

Differential Revision: [D125110873](https://our.internmc.facebook.com/intern/diff/D125110873/)
@Gasoonjia
Update
0ff7b05 
[ghstack-poisoned]
Gasoonjia added a commit that referenced this pull request Oct 19, 2026
@Gasoonjia
[executorch][cuda] Test CudaExecutor on every step width up to 65
0be3c22 
Pull Request resolved: #23819

CudaExecutor chooses the narrowest static forward_{N} that holds a step, pads unused rows into the KV scratch cell without returning their logits, and uses forward_others with slicing for wider steps. Extend the toy artifact to the Muse-sized method set forward_{1,2,4,...,64} plus forward_others, with eight independently referenced prompts and continuations.

EveryStepWidthUpTo65MatchesEagerGreedy covers every width 1..65 plus 96 and 97, shuffled repeatedly, with CUDA graphs both on and off. Mix sequence owners and chunk lengths, teacher-force eager reference histories, and vary which chunks request output. Assert the selected method sequence, exact output ownership/count and eager greedy token at each requested position. This is correctness coverage of static padding, dynamic widths and slicing, without serving performance reporting.
ghstack-source-id: 451855103
@exported-using-ghexport

Differential Revision: [D125110873](https://our.internmc.facebook.com/intern/diff/D125110873/)
meta-codesync Bot pushed a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][cuda] Autotune SDPA on a visible mask and tune pack GQA …
ff638be 
…with the tile

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #23735
* #23734
* #23974
* #23973
* #23972
* #23971
* #23970
* #23969
* #23968
* #23967
* #23966
* #23965
* #23964
* #23963
* #23962
* #23961
* #23960
* #23959
* #23958
* #23985
* #23984
* #23983
* #23982
* #23981
* #23980
* #23979
* #23978
* #23819
* #23818
* #23733
* #23732
* #23731
* #23730
* #23729
* __->__ #23727
* #23728

D123681196 declared representative KV lengths for the SDPA kernels that read
`KV_LEN_ptr`. A masked call still autotuned on the wrong work, for two reasons:

- **The mask.** Inductor fabricates bool inputs as all False, and the kernels
  skip every KV block no row of a tile can see. A masked call (e.g. the
  off-graph cell step, which always passes a visibility mask) therefore timed
  each config doing almost no attention and picked the tile that loads the
  mask fastest: BLOCK_M=64 / BLOCK_N=128 / 8 warps, a single 32-CTA wave for
  Muse Glimmer's 32-query step.
- **Pack GQA.** `sdpa()` decided it with a tile-utilization rule that assumed
  BLOCK_M 32 or 64 (FlashAttention's should_pack_gqa). For 32 queries in
  groups of 16 both utilizations are 1.0, so it did not pack: each of the 32
  query heads re-read its KV head's cells.

This change:
- declares `Mask_ptr` next to `KV_LEN_ptr` on `_sdpa_fwd_kernel`,
  `_sdpa_decode_splitk_kernel` and `_sdpa_small_query_splitk_kernel`: an
  all-visible mask per KV-length scenario, so each config is timed sweeping
  everything up to that length. Calls without a mask pass 0 and are untouched;
- makes `PACK_GQA` a config dimension of `_sdpa_fwd_kernel`: every tile is
  tuned packed and unpacked (only unpacked without GQA). The kernel derives
  its grid head count from `PACK_GQA`, and the grid reads it from the config.
  The non-power-of-2 kernel keeps the rule, now with its fixed BLOCK_M=32.

On the Muse Glimmer 30B cell step (T queries, 32 query / 2 KV heads, D=128,
65536-cell pool, A100), per layer at 4650 / 9300 live cells, relative to the
fastest config at that extent:

| T | before | this diff |
|---|---|---|
| 16 | 1.00x / 1.00x | 1.00x / 1.00x |
| 32 | 2.05x / 1.33x | 1.00x / 1.00x |
| 64 | 2.35x / 2.09x | 1.00x / 1.00x |

Differential Revision: D124908769

Pull Request resolved: #23727
meta-codesync Bot pushed a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][muse] Export the batching contract for CUDA with --use-o…
01a7b48 
…ffgraph-kv-cache

Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #23735
* #23734
* #23974
* #23973
* #23972
* #23971
* #23970
* #23969
* #23968
* #23967
* #23966
* #23965
* #23964
* #23963
* #23962
* #23961
* #23960
* #23959
* #23958
* #23985
* #23984
* #23983
* #23982
* #23981
* #23980
* #23979
* #23978
* #23819
* #23818
* #23733
* #23732
* #23731
* #23730
* __->__ #23729
* #23727
* #23728

`export_solo.py --use-offgraph-kv-cache` now means the same thing on both
backends: export the contract the batching runtime drives. MLX already wrote
its selected-logits decoder there; CUDA now writes the text-only artifact
CudaExecutor (backends/cuda/batching) runs, over one set of weights and a
cell-layout off-graph KV cache:

- `forward_{N}` for N in `--forward-sizes` (default 1, 2, 4, ..., 64): static,
  T = K = N, each captured into a CUDA graph at runtime. Their quantized linears
  trace the smallest QuantizedGemmFamily bucket that holds N.
- `forward_others`: dynamic, T in [max N, `--max-step-tokens`] (default 4096).
- `sample` / `sample_argmax` (extension/llm/batching/sampler.py) so logits stay
  on the device; the existing `--logits` returns them to the host instead.
- Metadata: selected logits, native context, cache geometry, min prefill chunk,
  and `get_offgraph_kv_max_cells` (`--max-cells`, default 65536).

Without the flag CUDA is unchanged: the in-graph `embed_text`,
`forward_from_embeddings` (prefill) and `decode_from_embedding` (decode).
`--mmproj` with the flag on CUDA is a usage error, since the batching artifact
is text-only.

Every forward method declares its step width at input_pos, delegate input 1
(D124428885): forward_others keeps the exported input order too, which the
CUDA backend's CheckOffGraphKVStepWidthPass checks.

This replaces the single-sequence off-graph CUDA export that only the solo
engine's runtime-owned-cache path consumed, and with it the
`solo-text-offgraph` CI mode that ran it through `solo_runner`. CUDA batching
gets its own CI mode later in this stack. The solo engine's now unreachable
off-graph plumbing is left for a separate cleanup.

The exporter is the landable form of the CUDA batching experiments (D123051484,
D124377337, D124377340, D124377341, D124428885), cell layout only; the batched-sequence
layout is out of scope.

Differential Revision: D124871956

Pull Request resolved: #23729
meta-codesync Bot pushed a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][muse] Hash prepared image identity without CommonCrypto
7379f1e 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #23735
* #23734
* #23974
* #23973
* #23972
* #23971
* #23970
* #23969
* #23968
* #23967
* #23966
* #23965
* #23964
* #23963
* #23962
* #23961
* #23960
* #23959
* #23958
* #23985
* #23984
* #23983
* #23982
* #23981
* #23980
* #23979
* #23978
* #23819
* #23818
* #23733
* #23732
* #23731
* __->__ #23730
* #23729
* #23727
* #23728

The MG prompt preparer keys an image's prefix identity with SHA-256 over its
canonical RGB, dimensions and vision grid. It called CommonCrypto, which only
exists on Apple platforms, so the batching preparation library could not build
for the CUDA backend on Linux.

A small FIPS 180-4 SHA-256 in prompt_preparation.cpp replaces it. The hashed
bytes and their order are unchanged, so keys are identical to before; the
existing fixed digest fixture in batching_runtime_test (canonical 56x28 image)
pins that.

Differential Revision: D124871958

Pull Request resolved: #23730
meta-codesync Bot pushed a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][muse] Run the MG batching runtime on CudaExecutor
aaa2d48 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #23735
* #23734
* #23974
* #23973
* #23972
* #23971
* #23970
* #23969
* #23968
* #23967
* #23966
* #23965
* #23964
* #23963
* #23962
* #23961
* #23960
* #23959
* #23958
* #23985
* #23984
* #23983
* #23982
* #23981
* #23980
* #23979
* #23978
* #23819
* #23818
* #23733
* #23732
* __->__ #23731
* #23730
* #23729
* #23727
* #23728

The batching runner and multiplexed worker now build and run on CUDA, through
the same shared stack as MLX: ServingRuntime (prefix cache, sessions, text
output) over batching::Runner and DecodeFirstScheduler, with the MG prompt
preparer. Only the executor differs.

- `runtime/cuda_backend_factory.{h,cpp}`: builds `CudaExecutor`
  (backends/cuda/batching) over the cell-layout artifact
  `export_solo.py --backend cuda --use-offgraph-kv-cache` writes, and the
  preparation spec from the program's metadata (vocabulary, context, widest
  step). The artifact is text only, so the spec has no vision and image
  requests are rejected during preparation.
- A zero `--max_session_tokens` shares the cell pool evenly between the
  physical sessions (requests, prefix snapshots and the capture lane), capped
  by the exported context, since a cell pool cannot reserve every session's
  full context.
- `create_muse_glimmer_backend` picks the executor the binary was built for;
  `--backend` defaults to it.
- CMake: `MUSE_GLIMMER_BUILD_BATCHING` / `_TESTS` accept a CUDA configuration,
  link `cuda_batching` instead of the MLX executor, link the preparation
  library to the CUDA runtime its vision path uses, and export dynamic symbols
  for the AOTI blob. The MLX configuration is unchanged.
- `make muse-glimmer-batching-cuda`, the counterpart of the MLX target.

Differential Revision: D124871959

Pull Request resolved: #23731
meta-codesync Bot pushed a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][cuda] Test prefix-cache reuse on CudaExecutor
a153301 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #23735
* #23734
* #23974
* #23973
* #23972
* #23971
* #23970
* #23969
* #23968
* #23967
* #23966
* #23965
* #23964
* #23963
* #23962
* #23961
* #23960
* #23959
* #23958
* #23985
* #23984
* #23983
* #23982
* #23981
* #23980
* #23979
* #23978
* #23819
* #23818
* #23733
* __->__ #23732
* #23731
* #23730
* #23729
* #23727
* #23728

The serving runtime's opt-in prefix cache (batching::PrefixCache, enabled by
`--prefix_cache_entries`) reuses KV through Session::clone_async, which on CUDA
reaches CudaExecutor::clone -> SessionTable::clone -> CellCache::seq_clone. The
cell layout clones by reference: the snapshot's id joins the owner bits of the
source's cells below the boundary, so reuse shares the prompt's KV instead of
copying it. Nothing needed to change to enable it; this adds the coverage that
shows it is correct on the GPU.

PrefixCacheResumesFromACapturedPrompt, on each toy artifact (dense, sparse,
device_sampling): every prompt's cold generation leaves a snapshot captured at
its prompt boundary. A lookup for the same prompt reuses all but its last token
and one for a prompt extending it reuses the whole snapshot; both resume on the
eager greedy continuation. So the clone carried the prompt's exact KV, and the
source's later decode writes, padding rows and the clones' own writes did not
disturb one another.

Differential Revision: D124871960

Pull Request resolved: #23732
meta-codesync Bot pushed a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][muse] Serve MG batching on CUDA from the launcher
8b9a21b 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #23735
* #23734
* #23974
* #23973
* #23972
* #23971
* #23970
* #23969
* #23968
* #23967
* #23966
* #23965
* #23964
* #23963
* #23962
* #23961
* #23960
* #23959
* #23958
* #23985
* #23984
* #23983
* #23982
* #23981
* #23980
* #23979
* #23978
* #23819
* #23818
* __->__ #23733
* #23732
* #23731
* #23730
* #23729
* #23727
* #23728

`serve_batching.py --backend cuda` launches the CUDA batching worker behind the
same FastAPI app, SessionRuntime, multiplexed worker client, chat template and
streaming parser as MLX. It requires `--data-path` (the AOTI weights blob) and
rejects `--pos-embed-path`, since the CUDA artifact is text only; image
requests fail during preparation.

BATCHING_SERVING.md gains a CUDA section: export with
`export_solo.py --backend cuda --use-offgraph-kv-cache`, build with
`make muse-glimmer-batching-cuda`, launch, and how resident sessions, prefix
snapshots and the capture lane share the exported cell pool.

Differential Revision: D124871961

Pull Request resolved: #23733
meta-codesync Bot pushed a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][ci] Test Muse Glimmer batching on CUDA in OSS CI
8ad11cf 
Stack from [ghstack](https://github.com/ezyang/ghstack/tree/0.15.0) (oldest at bottom):
* #23735
* #23734
* #23974
* #23973
* #23972
* #23971
* #23970
* #23969
* #23968
* #23967
* #23966
* #23965
* #23964
* #23963
* #23962
* #23961
* #23960
* #23959
* #23958
* #23985
* #23984
* #23983
* #23982
* #23981
* #23980
* #23979
* #23978
* #23819
* __->__ #23818
* #23733
* #23732
* #23731
* #23730
* #23729
* #23727
* #23728

Build and test Muse Glimmer CUDA batching in OSS CI using the shared serving APIs. The unit job builds the normal runner/worker and explicitly enables/builds the CI-only workload tool through build_muse_glimmer_batching_bench.sh, then runs preparation and export/dispatch tests plus CPU CI-tool regressions. The batching e2e mode explicitly builds the tool before invoking it, runs concurrent staggered requests with full completion budgets, and requires measured multi-session executor steps so serial execution cannot pass as batching coverage. Keep the shared server launcher smoke test and existing non-batching variants. CUDA workflow path filters and job gates include the relocated tooling and regression tests.

Differential Revision: D124871963

Pull Request resolved: #23818
@meta-codesync
meta-codesync Bot merged commit f068fa8 into gh/Gasoonjia/59/base Oct 20, 2026
281 of 289 checks passed
@meta-codesync
meta-codesync Bot deleted the gh/Gasoonjia/59/head branch October 20, 2026 01:23
@meta-codesync
meta-codesync Bot deployed to cherry-pick-bot October 20, 2026 01:23 Active
@Gasoonjia
Gasoonjia deployed to upload-benchmark-results October 20, 2026 03:13 — with GitHub Actions Active
@Gasoonjia Gasoonjia mentioned this pull request Oct 20, 2026
Merged
Gasoonjia added a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][cuda] Test CudaExecutor on every step width up to 65 (#2…
b1beef3 
…4011)

This PR was created by the merge bot to help merge the original PR into
the main branch.
ghstack PR number: #23819 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/59/base
ghstack PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/59/head
Merge bot PR base:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/58/orig
Merge bot PR head:
https://github.com/pytorch/executorch/tree/gh/Gasoonjia/59/orig
Differential Revision:
[D125110873](https://our.internmc.facebook.com/intern/diff/D125110873/)
@diff-train-skip-merge

This branch was successfully deployed

3 active deployments
upload-benchmark-results — 0ff7b05e Deployed Oct 20, 2026 by Gasoonjia via upload-benchmark-results #21093
cherry-pick-bot — 0ff7b05e Deployed Oct 20, 2026 by meta-codesync[bot] via Try to create a PR with ghstack /orig branch #1854
cadence — 0ff7b05e Deployed Oct 19, 2026 by Gasoonjia via hifi-op-test / hifi4 #34118
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@shoumikhin shoumikhin shoumikhin approved these changes

Assignees

No one assigned

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. meta-exported

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@Gasoonjia @shoumikhin