Sample logits on CUDA memory before copying out to CPU
Open
@larryliu0820 is already working on this.
Since Dec 9, 2025.
- Dominant language
- Python
- Stars
- 5k
- Forks
- 1.2k
- Avg merge
- 2d 10h
- Merged PRs (30d)
- 581
Description
🚀 The feature, motivation and pitch
For LLM/ASR models on ExecuTorch CUDA backend, we always copy logits to CPU memory, right after we exit CUDA delegate. Theoretically we can run sampler on CUDA and only copy the sampled result out. The other variant of this is to store the sampled result on CUDA until all the tokens are generated and dump all the results at once.
This requires all delegates and the sampler on the same CUDA stream, also some scratch memory allocated for sampler.
Alternatives
No response
Additional context
No response
RFC (Optional)
No response
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Assessment
This issue has not been assessed yet.