Skip to content

[executorch][muse] Copy only committed DFlash hidden rows to host - #23982

Open
Gasoonjia wants to merge 5 commits into
gh/Gasoonjia/81/basefrom
gh/Gasoonjia/81/head
Open

[executorch][muse] Copy only committed DFlash hidden rows to host#23982
Gasoonjia wants to merge 5 commits into
gh/Gasoonjia/81/basefrom
gh/Gasoonjia/81/head

Conversation

@Gasoonjia

@Gasoonjia Gasoonjia commented Oct 19, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

Snapshot the verification taps on device and copy only the accepted (committed)
rows to the host, instead of every verified row, mirroring D125629666 for the
batched CUDA path.

Differential Revision: D125705714

@Gasoonjia
Update
78b5b8d 
[ghstack-poisoned]
@pytorch-bot

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

Copy link
Copy Markdown

🔗 Helpful Links

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

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

⏳ 6 Pending, 3 Unrelated Failures

As of commit 698d4af with merge base b1beef3 (image):

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

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

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

@github-actions

github-actions Bot commented Oct 19, 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 mentioned this pull request Oct 19, 2026
Merged
@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 19, 2026
This was referenced Oct 19, 2026
Closed
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Open
Open
Open
Open
Open
Open
Open
@Gasoonjia
Update
182b254 
[ghstack-poisoned]
@Gasoonjia
Update
83839c4 
[ghstack-poisoned]
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 Bot pushed a commit that referenced this pull request Oct 20, 2026
@Gasoonjia
[executorch][cuda] Test CudaExecutor on every step width up to 65
f068fa8 
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
@Gasoonjia
Update
b283c8d 
[ghstack-poisoned]
@Gasoonjia
Update
698d4af 
[ghstack-poisoned]

This branch was successfully deployed

1 active deployment
cadence — 698d4af3 Deployed Oct 20, 2026 by Gasoonjia via hifi-op-test / hifi4 #34175
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