[executorch][muse] Add cuda dflash executor by Gasoonjia · Pull Request #23981 · pytorch/executorch

Gasoonjia · 2026-10-19T04:58:49Z

Stack from ghstack (oldest at bottom):

Add the CUDA DFlash batching executor, structured like the MLX one: one
batching::Executor that owns a target CudaExecutor and a draft cell cache and
runs drain / draft / speculate / rewind per batch. Owners of any context
length are co-batched; the target verifies every owner's candidates in one
forward. backend_factory dispatches --dflash on CUDA to it (replacing the
NotSupported guard), with the same MuseGlimmerBackendConfig and flags as MLX.

Only d=16 / n=15 is supported for now.

Differential Revision: D125705713

[ghstack-poisoned]

pytorch-bot · 2026-10-19T04:58:52Z

🔗 Helpful Links

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

📄 Preview Python docs built from this PR

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

⏳ 1 Pending, 4 Unrelated Failures

As of commit 81539bb with merge base b1beef3 ():

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

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

pull / nxp-mcuxpresso-test / linux-job (gh) (trunk failure)
RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
pull / unittest / windows / windows-job (gh) (trunk failure)
.ci/scripts/tests/test_muse_glimmer_batching_workload.py::MuseGlimmerBatchingWorkloadTest::test_workload_without_a_model
pull / unittest-editable / windows / windows-job (gh) (trunk failure)
.ci/scripts/tests/test_muse_glimmer_batching_workload.py::MuseGlimmerBatchingWorkloadTest::test_workload_without_a_model
pull / unittest-nxp-neutron / linux-job (gh) (trunk failure)
RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1

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

github-actions · 2026-10-19T04:59:35Z

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]

…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

…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

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

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

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

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

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

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 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 Pull Request resolved: #23819

[ghstack-poisoned]

Update

ee2219d

[ghstack-poisoned]

Gasoonjia requested a review from larryliu0820 as a code owner October 19, 2026 04:58

Gasoonjia deployed to cadence October 19, 2026 04:59 — with GitHub Actions Active

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

Gasoonjia deployed to cadence October 19, 2026 05:18 — with GitHub Actions Active

shoumikhin approved these changes Oct 19, 2026

View reviewed changes

Update

b1bd6a9

[ghstack-poisoned]

Gasoonjia deployed to cadence October 19, 2026 06:21 — with GitHub Actions Active

Gasoonjia deployed to cadence October 19, 2026 06:41 — with GitHub Actions Active

Update

4a8ae9e

[ghstack-poisoned]

Gasoonjia deployed to cadence October 19, 2026 08:15 — with GitHub Actions Active

Gasoonjia deployed to cadence October 19, 2026 08:39 — with GitHub Actions Active

Update

ca8f76e

[ghstack-poisoned]

Gasoonjia deployed to cadence October 20, 2026 03:47 — with GitHub Actions Active

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

Gasoonjia deployed to cadence October 20, 2026 04:11 — with GitHub Actions Active

Update

81539bb

[ghstack-poisoned]

Gasoonjia deployed to cadence October 20, 2026 08:46 — with GitHub Actions Active

Gasoonjia deployed to cadence October 20, 2026 09:06 — with GitHub Actions Active

Navigation Menu

Repository navigation

[executorch][muse] Add cuda dflash executor - #23981

[executorch][muse] Add cuda dflash executor#23981
Gasoonjia wants to merge 5 commits into
gh/Gasoonjia/80/basefrom
gh/Gasoonjia/80/head

Gasoonjia commented Oct 19, 2026 •

edited

Loading

Uh oh!

pytorch-bot Bot commented Oct 19, 2026 •

edited

Loading

Uh oh!

github-actions Bot commented Oct 19, 2026

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

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

Uh oh!

Uh oh!

pytorch-bot Bot commented Oct 19, 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/23981

⏳ 1 Pending, 4 Unrelated Failures

Uh oh!

github-actions Bot commented Oct 19, 2026

This PR needs a release notes: label

Uh oh!

This branch was successfully deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Gasoonjia commented Oct 19, 2026 •

edited

Loading

pytorch-bot Bot commented Oct 19, 2026 •

edited

Loading

This PR needs a `release notes:` label