Skip to content

[executorch][cuda] Let the CUDA delegate drive the off-graph KV step - #23294

Merged
Gasoonjia merged 4 commits into
mainfrom
gh/gasoonjia/191/orig
Oct 1, 2026
Merged

Gasoonjia merged 4 commits into
mainfrom
gh/gasoonjia/191/orig

Conversation

@pytorchbot

Copy link
Copy Markdown
Collaborator

This PR was created by the merge bot to help merge the original PR into the main branch.
ghstack PR number: #22989 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/191/base
ghstack PR head: https://github.com/pytorch/executorch/tree/gh/gasoonjia/191/head
Merge bot PR base: https://github.com/pytorch/executorch/tree/gh/gasoonjia/186/orig
Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/gasoonjia/191/orig
Differential Revision: D121131873
@diff-train-skip-merge

Pull Request resolved: #22815

Lower the neutral `kvcache::update_and_attend` contract to runtime-bound private tensors, retain their logical metadata without serializing payloads, and allow the weight loader to defer those bindings. Add validation for manifests and storage omission. Authored with Codex.
ghstack-source-id: 439161249
@exported-using-ghexport

Differential Revision: [D119571985](https://our.internmc.facebook.com/intern/diff/D119571985/)
Pull Request resolved: #22816

Own dense flat and ring KV allocations in the CUDA runtime, grow flat caches geometrically with stream-ordered preservation, rebind AOTI pointers, and expose reset and memory metrics for the single-session path. Derive allocation size, growth-copy pitch, metrics, and AOTI descriptors from a configurable storage dtype; BF16 remains the default. Authored with Codex.
ghstack-source-id: 439161268
@exported-using-ghexport

Differential Revision: [D119571986](https://our.internmc.facebook.com/intern/diff/D119571986/)
Pull Request resolved: #22989

The runner currently calls prepare_step/commit_step on CudaKVCache, which puts
a backend face in the runner's hands. The delegate is the natural owner: it
already runs on the host immediately before and after the compiled program, and
that is the only point where storage can grow and be rebound before the kernels
read it.

The one thing the delegate cannot work out for itself is how many tokens a step
writes. Lowering compiles the cache op away, so at runtime there is only
compiled code over pre-bound memory and nothing calls back out. The export side
knows the width from its own signature, so it declares where to find it as an
"input_index:dim" compile spec, and the delegate reads that input's extent --
shape metadata only, never contents, which are on device.

Committing from here also pairs the two halves on every exit: a run that fails
never advances the cache past tokens the kernels did not write.
ghstack-source-id: 439161280
@exported-using-ghexport

Differential Revision: [D121131873](https://our.internmc.facebook.com/intern/diff/D121131873/)
@pytorch-bot

pytorch-bot Bot commented Oct 1, 2026 •

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 67 Pending, 2 Unclassified Failures

As of commit cb538e9 with merge base 7cfcde7 (image):

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

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

@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 1, 2026
Base automatically changed from gh/gasoonjia/186/orig to main October 1, 2026 08:25
@github-actions

github-actions Bot commented Oct 1, 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 merged commit 605b561 into main Oct 1, 2026
278 of 280 checks passed
@Gasoonjia
Gasoonjia deleted the gh/gasoonjia/191/orig branch October 1, 2026 08:55

This branch was successfully deployed

1 active deployment
cadence — cb538e9a Deployed Oct 1, 2026 by Gasoonjia via hifi-op-test / hifi4 #30745
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants