Repository navigation
[executorch][muse] Export the batching contract for CUDA with --use-offgraph-kv-cache - #24005
Merged
Merged
Conversation
…ffgraph-kv-cache Pull Request resolved: #23729 `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. ghstack-source-id: 451855082 @exported-using-ghexport Differential Revision: [D124871956](https://our.internmc.facebook.com/intern/diff/D124871956/)
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/24005
Note: Links to docs will display an error until the docs builds have been completed. ⏳ No Failures, 381 PendingAs of commit 360baba with merge base 897d28f ( UNSTABLE - The following job is marked as unstable, possibly due to flakiness on trunk:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
meta-cla
Bot
added
the
CLA Signed
shoumikhin
approved these changes
Oct 20, 2026
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: #23729 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/45/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/45/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/43/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/Gasoonjia/45/orig
Differential Revision: D124871956
@diff-train-skip-merge