Repository navigation
[executorch][muse] Run the CUDA DFlash drafter in CUDA graphs - #24043
Merged
Merged
Conversation
Pull Request resolved: #23985 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. ghstack-source-id: 453099299 @exported-using-ghexport Differential Revision: [D125705719](https://our.internmc.facebook.com/intern/diff/D125705719/)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24043
Note: Links to docs will display an error until the docs builds have been completed. This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
Gasoonjia
added this pull request to stack #24044
October 21, 2026 06:44
Gasoonjia
approved these changes
Oct 21, 2026
Gasoonjia
added a commit
that referenced
this pull request
Oct 21, 2026
CudaDFlashNative.FinalWholeBlockFitsExactlyAndPreservesAllPrefillRows has failed in unittest-cuda-muse-glimmer-batching since #24043 added draft forward count checks to it: Expected: (drained.draft_forwards - drain.draft_forwards) > (drained.static_draft_forwards - drain.static_draft_forwards), actual: 1 vs 1 The test expected the decode at position 78 to drain the 79 hidden rows dynamically before the static draft. But run_target drains every owner it touches right after appending hidden rows (#24042), so the 78-row prefill already drains down to the fresh context (32). The decode then holds at most one block of rows and takes a single static draft, and the dynamic forwards fall before the window the test measured. Measure from before the prefill instead: the prefill drains through dynamic drafts only, and the decode runs exactly one static draft.
Gasoonjia
added a commit
that referenced
this pull request
Oct 22, 2026
Stacked on #24047. `CudaDFlashNative.FinalWholeBlockFitsExactlyAndPreservesAllPrefillRows` has failed in `unittest-cuda-muse-glimmer-batching` since #24043. It failed on that PR's own CI and on `main` at 8d9eec3: ``` cuda_dflash_native_test.cpp:612: Failure Expected: (drained.draft_forwards - drain.draft_forwards) > (drained.static_draft_forwards - drain.static_draft_forwards), actual: 1 vs 1 ``` The test expected the decode at position 78 to drain the 79 hidden rows through the dynamic draft before running the static draft. But `run_target` drains every owner it touches right after it appends hidden rows (added in #24042). So the 78-row prefill already drains down to the fresh context (32 in this fixture). The decode then holds at most one block (16) of rows and runs a single static draft. The dynamic forwards happen during the prefill, before the stats snapshot the test took. This is a bug in the test's measurement window, not in the executor. The fix takes the first snapshot before the prefill and checks two things: - the prefill runs only dynamic draft forwards (at least one); - the decode runs exactly one draft forward, and it is static. Test plan: - `clang-format` is clean on the changed lines. - No local GPU, so the test relies on CI: `unittest-cuda-muse-glimmer-batching` on this PR should pass.
This branch was successfully deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #23985 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/84/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/84/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/83/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/84/orig
Differential Revision: D125705719
@diff-train-skip-merge