Skip to content

perf(v4): batch decode sampling into a single sync (+3.6-11.7% tps) - #88

Merged
nosyndicate merged 1 commit into
mainfrom
phase4_40
Jul 23, 2026
Merged

perf(v4): batch decode sampling into a single sync (+3.6-11.7% tps)#88
nosyndicate merged 1 commit into
mainfrom
phase4_40

Conversation

@nosyndicate

@nosyndicate nosyndicate commented Jul 22, 2026

Copy link
Copy Markdown
Owner

Replaces v4's per-row decode sampling loop (up to batch host↔device .item() syncs per step) with a single batched sample_tokens call plus one .tolist() sync. Pure host-side optimization confined to ScheduleInferenceEngine (v4-only); randomness stays deterministic and batch-invariant via uniforms_from_seeds, keyed on each request's (seed, output-token count).

  • Adds 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
  • Refactors token-finalization (EOS/detok/max-len) into _decode_result_for_token, shared by the prefill path (_sample_one) and the new batched decode path (_post_decode)
  • _post_decode now gathers per-row temperature/top_k/top_p into tensors, draws deterministic uniforms for the whole batch, and calls sample_tokens once instead of looping sample_token per sequence
  • Adds test_post_decode_mixed_greedy_sampled_topk_batch, covering a batch mixing greedy, top-k=1, and top-k=2 rows

Benchmarked on RTX A6000 / Qwen3-1.7B (v4 paged attention), batch 8→64, concurrency = batch (no starvation confound):

batch Δ throughput Δ TPOT
8 +3.6% −3.7%
16 +7.0% −6.8%
32 +6.8% −6.2%
64 +11.7% −10.8%

Peak GPU memory identical at every batch; pytest tests/executor 135 passed (main: 134). Benefit scales with batch size (more eliminated per-row syncs).

@nosyndicate nosyndicate changed the title using batch sampling perf(v4): batch decode sampling into a single sync (+3.6-11.7% tps) Jul 23, 2026
@nosyndicate
nosyndicate marked this pull request as ready for review July 23, 2026 08:33
Copilot AI review requested due to automatic review settings July 23, 2026 08:33
@nosyndicate
nosyndicate merged commit fe6d755 into main Jul 23, 2026
1 check passed
@nosyndicate
nosyndicate deleted the phase4_40 branch July 23, 2026 08:37

Copilot AI left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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_decode using per-row parameter tensors + uniforms_from_seeds, replacing the prior per-row sampling loop.
  • Introduce GenerationRequestState.noise_salt and 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_token and 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 thread server/executor/engine.py
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 thread server/executor/engine.py
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)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants