[executorch][muse] Draft CUDA DFlash decode from pending tokens

[Gasoonjia]

Stack from ghstack (oldest at bottom):

• [executorch][ci] Test Muse Glimmer batching on CUDA in OSS CI #23735
• [executorch][muse] Benchmark MG batching through the serving runtime #23734
• [executorch][muse] Serve Muse Glimmer CUDA batching on a batched-sequence cache #23974
• [executorch][cuda] Lower off-graph KV to the batched-sequence layout #23973
• [executorch][cuda] Add a batched-sequence CUDA KV cache #23972
• [executorch][cuda] Add a length-aware Triton SDPA over per-sequence KV columns #23971
• [executorch][muse] Accept images in CUDA DFlash batching #23970
• [executorch][muse] Serve images on CUDA batching #23969
• [executorch][muse] Export vision with the CUDA batching contract #23968
• [executorch][cuda] Run CudaExecutor steps from prepared embeddings #23967
• [executorch][muse] Share the Muse Glimmer embeddings engine across backends #23966
• [executorch][muse] Run the CUDA DFlash drafter in CUDA graphs #23965
• [executorch][muse] Verify CUDA DFlash drafts on device under any sampling #23964
• [executorch][muse] Draft CUDA DFlash decode from pending tokens #23963
• [executorch][muse] Copy only committed DFlash hidden rows to host #23962
• [executorch][muse] Add cuda dflash executor #23961
• [executorch][muse] Export packed DFlash drafts for CUDA batching #23960
• [executorch][cuda] Expose raw packed steps on CudaExecutor #23959
• [executorch][cuda] Run packed noncausal DFlash draft steps on the cell cache #23958
• [executorch][muse] Run the CUDA DFlash drafter in CUDA graphs #23985
• [executorch][muse] Verify CUDA DFlash drafts on device under any sampling #23984
• -> [executorch][muse] Draft CUDA DFlash decode from pending tokens #23983
• [executorch][muse] Copy only committed DFlash hidden rows to host #23982
• [executorch][muse] Add cuda dflash executor #23981
• [executorch][muse] Export packed DFlash drafts for CUDA batching #23980
• [executorch][cuda] Expose raw packed steps on CudaExecutor #23979
• [executorch][cuda] Run packed noncausal DFlash draft steps on the cell cache #23978
• [executorch][cuda] Test CudaExecutor on every step width up to 65 #23819
• [executorch][ci] Test Muse Glimmer batching on CUDA in OSS CI #23818
• [executorch][muse] Serve MG batching on CUDA from the launcher #23733
• [executorch][cuda] Test prefix-cache reuse on CudaExecutor #23732
• [executorch][muse] Run the MG batching runtime on CudaExecutor #23731
• [executorch][muse] Hash prepared image identity without CommonCrypto #23730
• [executorch][muse] Export the batching contract for CUDA with --use-offgraph-kv-cache #23729
• [executorch][cuda] Retain emptied KV backing across JIT autotuning #23728
• [executorch][cuda] Autotune SDPA on a visible mask and tune pack GQA with the tile #23727

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

[pytorch-bot]

🔗 Helpful Links

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

• 📄 Preview Python docs built from this PR

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

❌ 3 New Failures, 1 Pending, 1 Unclassified Failure

As of commit ddb1bd5 with merge base 897d28f ():

NEW FAILURES - The following jobs have failed:

• Build Presets / apple (ios-simulator) / build (gh)
  ##[error]fatal: unable to access 'https://github.com/facebook/buck2-prelude.git/': Failed to connect to github.com port 443 after 603 ms: Couldn't connect to server
• pull / unittest / windows / windows-job (gh)
  .ci/scripts/tests/test_muse_glimmer_batching_workload.py::MuseGlimmerBatchingWorkloadTest::test_workload_without_a_model
• pull / unittest-editable / windows / windows-job (gh)
  .ci/scripts/tests/test_muse_glimmer_batching_workload.py::MuseGlimmerBatchingWorkloadTest::test_workload_without_a_model

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:

• Test CUDA Builds / test-muse-glimmer-cuda-e2e-kquant-17gb-solo-text-batching / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  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]

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.