perf(v4): batch decode sampling into a single sync (+3.6-11.7% tps) - #88
Merged
Conversation
nosyndicate
marked this pull request as ready for review
July 23, 2026 08:33
There was a problem hiding this comment.
Pull request overview
This PR optimizes v4 ScheduleInferenceEngine decode-time sampling by batching per-sequence sampling into a single sample_tokens call, reducing host↔device synchronization overhead while keeping sampling deterministic and batch-invariant via uniforms_from_seeds.
Changes:
- Batch decode sampling in
_post_decodeusing per-row parameter tensors +uniforms_from_seeds, replacing the prior per-row sampling loop. - Introduce
GenerationRequestState.noise_saltand engine-side_noise_seed()to ensure unseeded requests still have a stable per-request seed for counter-based randomness. - Refactor token finalization into
_decode_result_for_tokenand add a targeted unit test for mixed greedy/top-k rows in a single batched sampling step.
Reviewed changes
Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.
| File | Description |
|---|---|
server/executor/engine.py |
Implements batched decode sampling (sample_tokens), adds per-request noise seeding, and centralizes token finalization logic. |
server/executor/types.py |
Adds noise_salt to request state to support deterministic counter-based sampling for unseeded requests. |
tests/executor/test_schedule_engine.py |
Adds coverage for batched _post_decode with mixed greedy and top-k sampling rows. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+385
to
+390
| seed = request_state.sampling_params.seed | ||
| if seed is not None: | ||
| return seed | ||
| if request_state.noise_salt is None: | ||
| request_state.noise_salt = random.getrandbits(63) | ||
| return request_state.noise_salt |
Comment on lines
654
to
+657
| token_id = sample_token( | ||
| logits, request_state.sampling_params, request_state.generator | ||
| ) | ||
| return self._decode_result_for_token(token_id, request_state) |
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.
Replaces v4's per-row decode sampling loop (up to
batchhost↔device.item()syncs per step) with a single batchedsample_tokenscall plus one.tolist()sync. Pure host-side optimization confined toScheduleInferenceEngine(v4-only); randomness stays deterministic and batch-invariant viauniforms_from_seeds, keyed on each request's(seed, output-token count).noise_salt: a stable per-request random seed assigned at admission for unseeded requests, so the counter-based batched noise source always has a concrete integer seed_decode_result_for_token, shared by the prefill path (_sample_one) and the new batched decode path (_post_decode)_post_decodenow gathers per-row temperature/top_k/top_p into tensors, draws deterministic uniforms for the whole batch, and callssample_tokensonce instead of loopingsample_tokenper sequencetest_post_decode_mixed_greedy_sampled_topk_batch, covering a batch mixing greedy, top-k=1, and top-k=2 rowsBenchmarked on RTX A6000 / Qwen3-1.7B (v4 paged attention), batch 8→64, concurrency = batch (no starvation confound):
Peak GPU memory identical at every batch;
pytest tests/executor135 passed (main: 134). Benefit scales with batch size (more eliminated per-row syncs).