Repository navigation
feat(inference): support request-local sampling seeds - #7828
Merged
Merged
Conversation
Signed-off-by: wdykas <wdykas@nvidia.com>
Contributor
Author
|
/review model=claude |
wdykas
commented
Oct 6, 2026
Contributor
Author
|
/review |
Contributor
Author
|
this has a scalability problem I need to figure out |
Replace per-request CUDA generator reseeding and exponential launches with a batched Philox kernel keyed by request seed, absolute token position, and vocabulary index. Preserve the CPU fallback and unseeded shared-generator behavior. Document the seeded CUDA stream change, cover replay across batch layouts and launch tiling under contention, and add a reproducible sampler benchmark. Validated with 333 focused tests on GB200. Signed-off-by: wdykas <wdykas@nvidia.com>
Keep the PR focused on production request-local sampling and necessary regression coverage. Remove the ad hoc benchmark, duplicate CUDA replay wrapper, and sampler position-override test already covered by the async-controller regression. Retain the required kernel determinism and distribution checks. Validated with 331 focused tests and formatting, lint, and kernel coverage checks. Signed-off-by: wdykas <wdykas@nvidia.com>
Contributor
Author
|
/ok to test f2bb8b7 |
Signed-off-by: wdykas <wdykas@nvidia.com>
Signed-off-by: wdykas <wdykas@nvidia.com>
wdykas
marked this pull request as ready for review
October 6, 2026 16:36
Contributor
Author
|
/ok to test cdd420a |
wdykas
commented
Oct 7, 2026
|
|
||
|
|
||
| @triton.jit(do_not_specialize=["Seeds", "Positions"]) | ||
| def _request_seed_noise_kernel( |
Contributor
Author
There was a problem hiding this comment.
needed for perf. the naive way takes forever
Contributor
Author
|
/ok to test ca4e67c |
Signed-off-by: wdykas <wdykas@nvidia.com>
Contributor
Author
|
/ok to test e5a35e3 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38096793901 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38100823807 |
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Oct 11, 2026
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38145549091 |
github-merge-queue
Bot
removed this pull request from the merge queue due to failed status checks
Oct 11, 2026
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38161303560 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/38165752174 |
This was referenced Oct 14, 2026
Waterpine
added a commit
to Waterpine/Megatron-LM
that referenced
this pull request
Oct 20, 2026
torch_sampling.py is registered under the request-seed sampling kernel since NVIDIA#7828, so changes to it must update that kernel's determinism test. Replay an all-greedy batch with request seeds under side-stream contention, and check that it returns per-row argmax without advancing the shared RNG. Also update the sample_kernel docstring, which still described no_top_k/no_top_p as ignored. Signed-off-by: Waterpine <biansonghz@gmail.com>
2 of 7 tasks
4 of 7 tasks
This branch was previously 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.
What does this PR do?
Add optional request-local sampling seeds to dynamic Torch inference and both OpenAI-compatible completion endpoints. With a shared generator, moving a request between batches changes its draws even when its logits are unchanged. Explicit
seedassigns each draw by request seed and absolute token position, independently of request IDs, batch slots, or other requests' RNG consumption.iseed(seed + i) % 2**63, including streaming submissions.Issue tracking
Closes #7827