Skip to content

[executorch][muse] Benchmark MG batching through the serving runtime - #23734

Open
Gasoonjia wants to merge 1 commit into
gh/Gasoonjia/50/basefrom
gh/Gasoonjia/50/head
Open

[executorch][muse] Benchmark MG batching through the serving runtime#23734
Gasoonjia wants to merge 1 commit into
gh/Gasoonjia/50/basefrom
gh/Gasoonjia/50/head

Conversation

@Gasoonjia

@Gasoonjia Gasoonjia commented Oct 13, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Stack from ghstack (oldest at bottom):

muse_glimmer_batching_bench measures throughput where requests really enter:
ServingRuntime, its Runner and DecodeFirstScheduler, over the backend
executor (CUDA or MLX), with the MG preparer and text output. It ports
run_solo_batching's measurements from the CUDA batching experiments onto this
stack instead of driving the executor directly.

--mode=sweep keeps every step at T tokens for each T in --tokens (default
1, 2, 4, ..., 4096), the step mix of the experiments' T sweep:

  • T <= 64: T requests decoding together;
  • wider: --decode_cap (32) decodes plus a backlog of prompts that fill the
    rest of each step. By default one prompt fills it: ServingRuntime opens
    about one new session per engine step (its control thread waits on each
    open), so several fresh prompts per step cannot be admitted fast enough to
    form the shape. --prefill_chunk N splits the prefill into N-token prompts.
    A thin Executor decorator records each step's width and time. Each row is the
    median over --iters steady steps after --warmup, reported both as executor
    time and as the interval between consecutive steps; the latter adds the
    runtime's own per-step work (scheduling, callbacks, detokenization) and is the
    headline tok/s. The executor loads once; each T gets a fresh ServingRuntime.

--mode=requests: --requests prompts of --prompt_tokens tokens submitted
--stagger_ms apart, each generating --max_new_tokens; reports aggregate
tok/s and each request's TTFT and TPOT.

Differential Revision: D124871962

@Gasoonjia
Update
b21aa62 
[ghstack-poisoned]
@pytorch-bot

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

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 6 New Failures, 1 Cancelled Job

As of commit b21aa62 with merge base c508445 (image):

NEW FAILURES - The following jobs have failed:

CANCELLED JOB - The following job was cancelled. Please retry:

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

This was referenced Oct 13, 2026
Merged
Closed
Merged
Merged
Merged
Merged
Merged
@Gasoonjia Gasoonjia mentioned this pull request Oct 13, 2026
Open
@github-actions

github-actions Bot commented Oct 13, 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.

@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 13, 2026
This was referenced Oct 14, 2026
Merged
Merged
Open
Open
Open
Open
Open
Open
Open
Open
Open
Open
This was referenced Oct 18, 2026
Open
Open
Open
Open
Open
Open
Merged
Merged
Merged
Merged
Merged
Merged
Merged
Merged
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
meta-codesync Bot pushed a commit that referenced this pull request Oct 21, 2026
@Gasoonjia
[executorch][cuda] Run packed noncausal DFlash draft steps on the cel…
099778c 
…l 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

MLX batches DFlash drafts by packing every owner's new target context and
query block into one noncausal draft step. Let the CUDA cell cache serve the
same step so the CUDA DFlash executor can reuse the off-graph KV path.

- CacheGeometry::dflash_block_size (already on main) now selects a packed
  draft layout on CUDA: context rows commit, query rows are scratch.
- CudaCellCache prepares packed draft steps (declared context + block rows per
  owner) and lowers the noncausal mask the draft attends with.
- lower_offgraph_kv / cuda_backend learn the draft step width spec.
- Drop the "DFlash is unsupported on CUDA" builder guard.

Differential Revision: D125705710

Pull Request resolved: #23978
meta-codesync Bot pushed a commit that referenced this pull request Oct 21, 2026
@Gasoonjia
[executorch][cuda] Expose raw packed steps on CudaExecutor
5a67024 
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

DFlash needs the target's verify rows and tapped hidden states from one
forward without the executor's own sampling. Add a RawTargetStep /
execute_raw API and argmax_rows to CudaExecutor, reusing its batch planning,
step declaration and commit, so a model adapter can drive verification the way
the MLX DFlash executor drives its target.

Differential Revision: D125705711

Pull Request resolved: #23979
meta-codesync Bot pushed a commit that referenced this pull request Oct 21, 2026
@Gasoonjia
[executorch][muse] Export packed DFlash drafts for CUDA batching
85d220f 
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

Export a tapped target plus a packed draft_forward for CUDA batching, matching
the MLX DFlash draft ABI: tokens[owners, block], hidden[1, C, H],
positions[C + Q], order[C + Q] -> logits[1, Q, V]. block = n_draft + 1.

- export_dflash.py: --use-offgraph-kv-cache with --backend cuda emits the
  packed draft over a draft cell cache (--draft-max-cells, --max-draft-owners).
- export_solo.py: device residency and cache wiring for draft methods.
- dflash_model.py: packed draft forward; cuda.py: hidden tapping for CUDA.

Differential Revision: D125705712

Pull Request resolved: #23980
meta-codesync Bot pushed a commit that referenced this pull request Oct 21, 2026
@Gasoonjia
[executorch][muse] Add cuda dflash executor
6f1bc1f 
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

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

Pull Request resolved: #23981
meta-codesync Bot pushed a commit that referenced this pull request Oct 21, 2026
@Gasoonjia
[executorch][muse] Copy only committed DFlash hidden rows to host
637ea35 
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

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

Pull Request resolved: #23982
meta-codesync Bot pushed a commit that referenced this pull request Oct 21, 2026
@Gasoonjia
[executorch][muse] Draft CUDA DFlash decode from pending tokens
109a09b 
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

After a verify, the emitted correction/bonus token is kept as pending instead
of running a target preamble for it; the next decode drafts directly from it.
This removes one target forward per speculation round, as on MLX.

Differential Revision: D125705715

Pull Request resolved: #23983
meta-codesync Bot pushed a commit that referenced this pull request Oct 21, 2026
@Gasoonjia
[executorch][muse] Verify CUDA DFlash drafts on device under any samp…
0018036 
…ling

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

Export dflash_verify / dflash_verify_argmax: batched on-device acceptance of
every owner's drafts under that owner's temperature / top-k / top-p, with a
greedy fast path. Only token IDs return to the host. Uniforms come from
(seed, position, purpose) so results do not depend on batch composition.
Target preamble sampling also stays on device (CudaExecutor::sample_rows).

Differential Revision: D125705716

Pull Request resolved: #23984
meta-codesync Bot pushed a commit that referenced this pull request Oct 21, 2026
@Gasoonjia
[executorch][muse] Run the CUDA DFlash drafter in CUDA graphs
4feed35 
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

Export static draft_forward_{N} (N owners, at most one block of new context
each) and capture them in CUDA graphs through the existing CudaExecutor graph
option; longer contexts still drain through the dynamic draft_forward. Decode
contexts are padded to the static shape with scratch rows the cache never
commits. Pool growth recaptures like the target's static methods.

Differential Revision: D125705719

Pull Request resolved: #23985

This branch was successfully deployed

1 active deployment
cadence — b21aa62c Deployed Oct 13, 2026 by Gasoonjia
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@larryliu0820 larryliu0820 Awaiting requested review from larryliu0820 larryliu0820 is a code owner
@kirklandsign kirklandsign Awaiting requested review from kirklandsign

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.

1 participant

@Gasoonjia