-
Notifications
You must be signed in to change notification settings - Fork 779
Add CudaSampler class for GPU-based token sampling #16387
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
base: gh/larryliu0820/87/head
Are you sure you want to change the base?
Conversation
🔗 Helpful Links🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/16387
Note: Links to docs will display an error until the docs builds have been completed. ❌ 3 New Failures, 3 Unrelated FailuresAs of commit 8573c9a with merge base c5d66a5 ( NEW FAILURES - The following jobs have failed:
BROKEN TRUNK - The following jobs failed but were present on the merge base:👉 Rebase onto the `viable/strict` branch to avoid these failures
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. |
| // - The argmax kernel will wait for the decoder to finish writing logits | ||
| // - No explicit cudaDeviceSynchronize() or cross-stream synchronization needed | ||
| // | ||
| // 4. Trade-off: Using the default stream prevents concurrent execution between |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Im wondering if we really need to make sampler and cuda backend using same cuda stream, since the sampling and decoding should be able to work in parallel: the argmax process of logits_{i} should be able to work with the decoder generating logits_{i+1} since they do not have any dependency, and such parallelism may not happen if argmax and decoder share the same cudastream.
Add CudaSampler class that provides a high-level interface for GPU sampling:
Pre-allocates GPU memory to avoid allocation in hot path.
for implicit synchronization with the CUDA backend's stream.
The default stream approach ensures proper ordering between decoder
output and argmax without requiring explicit cross-stream synchronization
or access to the backend's internal stream.