feat(serving): support async scheduling for DeepSeek MTP decode - #130
Open
ndleslx wants to merge 2 commits into
Open
feat(serving): support async scheduling for DeepSeek MTP decode#130ndleslx wants to merge 2 commits into
ndleslx wants to merge 2 commits into
Conversation
|
Important Review skippedAuto incremental reviews are disabled on this repository. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Organization UI Review profile: CHILL Plan: Pro Plus Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
📝 WalkthroughWalkthroughAsync scheduling now defaults to enabled. Async decode scheduling reserves capacity for speculative tokens and reconciles placeholders and computed-token counts when tokens are accepted, rejected, or terminated by EOS. ChangesAsync speculative scheduling
Estimated code review effort: 3 (Moderate) | ~25 minutes Sequence Diagram(s)sequenceDiagram
participant Scheduler
participant Request
participant Worker
Scheduler->>Request: Reserve base and speculative placeholders
Scheduler->>Worker: Schedule async decode
Worker-->>Scheduler: Return sampled tokens
Scheduler->>Request: Retain tokens and reconcile placeholders
Scheduler->>Request: Correct computed-token count for shortfall or EOS
Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
🧹 Nitpick comments (2)
pypto_serving/serving/sched/scheduler.py (2)
566-591: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winMove block-publish after the shortfall correction for robustness.
_cache_completed_blocks(request)runs beforenum_computed_tokensis corrected for the shortfall below. This is harmless today only becauserequest.block_hashesis fixed at prompt-computation time (add_request,compute_block_hashes(request.prompt_token_ids)), sototal_blocks_computedis already capped at the prompt boundary by the time any decode-step shortfall can occur — the over-counted value never crosses into new cache entries. If prefix caching is ever extended to cover generated tokens, this ordering would publish blocks for speculative positions that were actually rejected/dropped. Reordering costs nothing and removes the implicit dependency.🛡️ Proposed reorder
- # Publish confirmed blocks now (worker succeeded), not at dispatch. - self._cache_completed_blocks(request) - retained_tokens = 0 for token_id in token_ids: request.output_token_ids.append(token_id) retained_tokens += 1 outputs.append(RequestOutput(request_id=request.request_id, new_token_id=token_id)) if self._check_finish(request) is not None: # Tokens after a finish are dropped (mirrors the sync path), so # they must not count as retained. break if reserved: # Release every placeholder this step reserved. request.num_output_placeholders = max( 0, request.num_output_placeholders - reserved ) # Give back the optimistically-advanced positions that produced no # retained token (speculative rejection, or an early EOS). shortfall = reserved - retained_tokens if shortfall > 0: request.num_computed_tokens = max( 0, request.num_computed_tokens - shortfall ) + + # Publish confirmed blocks now (worker succeeded), not at dispatch. + # Deferred until after the shortfall correction so a future caching + # extension can't publish blocks for un-retained speculative tokens. + self._cache_completed_blocks(request)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/sched/scheduler.py` around lines 566 - 591, Move the _cache_completed_blocks(request) call in the decode completion flow to after the retained-token loop and the reserved-placeholder shortfall correction. Ensure num_computed_tokens is corrected before completed blocks are published, while preserving the existing token handling and placeholder release behavior.
388-425: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDocstring no longer matches the reservation logic.
The docstring still says decode steps "reserve one
num_output_placeholders," but the code below now reserves1 + self._speculative_tokens_for(...)(up to1 + num_speculative_tokens) for greedy decode/MTP steps. The sibling docstring in_reconcile_async_output(Lines 547-552) was updated to describe this upper-bound model; this one was missed.📝 Proposed docstring update
- - decode requests reserve one ``num_output_placeholders`` for the token - this step will sample but that is not yet known. Prefill chunks that do - not complete the prompt sample nothing, so they add no placeholder; - the chunk that completes the prompt reserves one (its first generated - token), matching the sync path where that token is appended. + - sampling steps (decode, or the prefill chunk that completes the + prompt) reserve ``1 + speculative_tokens`` ``num_output_placeholders`` + — the upper bound of tokens this step may emit, since the accepted + count is not yet known for speculative/MTP decode. Prefill chunks + that do not complete the prompt sample nothing, so they add none.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@pypto_serving/serving/sched/scheduler.py` around lines 388 - 425, The docstring for the async scheduling reservation logic should describe that decode steps reserve the upper bound of 1 + self._speculative_tokens_for(request, scheduled) output placeholders, rather than exactly one. Update the decode/MTP explanation to mention the variable result range and that reconciliation removes any unused reservation; keep the prefill completion behavior unchanged.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Nitpick comments:
In `@pypto_serving/serving/sched/scheduler.py`:
- Around line 566-591: Move the _cache_completed_blocks(request) call in the
decode completion flow to after the retained-token loop and the
reserved-placeholder shortfall correction. Ensure num_computed_tokens is
corrected before completed blocks are published, while preserving the existing
token handling and placeholder release behavior.
- Around line 388-425: The docstring for the async scheduling reservation logic
should describe that decode steps reserve the upper bound of 1 +
self._speculative_tokens_for(request, scheduled) output placeholders, rather
than exactly one. Update the decode/MTP explanation to mention the variable
result range and that reconciliation removes any unused reservation; keep the
prefill completion behavior unchanged.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro Plus
Run ID: ba4520e6-d1da-4056-97f1-2e298004a092
📒 Files selected for processing (3)
pypto_serving/serving/engine/async_engine.pypypto_serving/serving/sched/scheduler.pytests/test_batching.py
ndleslx
force-pushed
the
serving-async-mtp
branch
3 times, most recently
from
July 31, 2026 08:43
df749ca to
992f7d2
Compare
Async (depth-2 pipelined) scheduling was gated off for the MTP executor because a speculative decode step returns a variable number of accepted tokens (1 .. 1+num_speculative_tokens), only known once the worker replies, while the optimistic advance assumed exactly one token per step. Enable it by reserving the upper bound at dispatch and correcting on reconciliation. Scheduler: - advance_after_schedule() now reserves 1 + num_speculative_tokens placeholders for a sampling step and advances num_computed_tokens by the same upper bound. This matches the blocks schedule() already allocates (num_new + speculative_tokens), so no extra KV pressure is introduced. - _speculative_tokens_for() mirrors schedule()'s rule: speculative capacity applies only to greedy (temperature <= 0) decode steps, never prefill. - _reconcile_async_output() releases every placeholder the step reserved and subtracts the shortfall (reserved - retained) from num_computed_tokens, so a rejected draft — or an early EOS that drops the trailing token — lands on exactly the same state as the synchronous path. Engine: - resolve_async_scheduling() no longer excludes PyptoDeepSeekV4Executor; async now defaults on for all executors. Explicit async_scheduling=False is still honoured for an A/B or a fallback. No worker changes were required: DeepSeek resolves prev_hidden_states on the host from the prev token id, and placeholder substitution (_resolve_decode_token / _resolve_prev_token) already runs before that lookup, so MTP's non-None prev-token requirement is satisfied. The worker's _last_tokens cache already records every returned token and keeps the last two, which is exactly the last/prev context MTP needs. MTP's per-request draft state (_mtp_request_states) is written inside the decode call, and the worker is strictly serial, so pipelining never overlaps two worker steps. Tests: upper-bound reservation; async == sync with both tokens accepted; async == sync when the draft is rejected (no permanent desync); EOS mid-pair drops the trailing token and reclaims its position.
The serving layer computed a prev-token per decode request and, for executors without device embedding, an extra [batch, hidden] embedding tensor for it — all of which the model layer discards. DeepSeekV4's MTP flow overwrites both fields before the kernel sees them (_main_speculative_batch): prev_token_ids is rebuilt from batch.token_ids (the last token) and prev_hidden_states from batch.hidden_states, so the values the worker supplied are unused. The only reader, prepare_mtp_decode_inputs, runs on that already-overwritten speculative batch. DeepSeekV4's autoregressive flow and the Qwen runner never read the fields at all. Remove the dead plumbing: - DecodeRequest.prev_token (IPC) and the engine code that derived it. - WorkerProcess._resolve_prev_token and the per-step prev token tensor plus its lookup_embeddings() call, which was wasted host work on every decode step. DecodeBatch.prev_token_ids / prev_hidden_states stay: they remain the interface the MTP runner populates for its own verification pair. The MTP speculative-batch tests, which assert the overwritten prev_token_ids and prev_hidden_states, still pass — confirming the removed values never reached the kernel.
ndleslx
force-pushed
the
serving-async-mtp
branch
from
July 31, 2026 09:19
992f7d2 to
04ca067
Compare
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
1 participant
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.
Enables async (depth-2 pipelined) scheduling for the DeepSeek V4 MTP decoder and removes dead prev-token plumbing found along the way.
Both on-device accuracy guards pass: Qwen 4/4, DeepSeek 4/4.
Two commits
1.
feat(serving): support async scheduling for DeepSeek MTP decodeScheduler — optimistic reservation + reconciliation:
advance_after_schedule()reserves1 + num_speculative_tokensplaceholders and advancesnum_computed_tokensby the speculative extra, matching the blocksschedule()already allocates._reconcile_async_output()releases placeholders and reclaims only unused speculative positions; a step's ownnum_new_tokensis never reverted (fixes a completing-prefill re-scheduling bug)._speculative_tokens_for()mirrorsschedule()'s greedy-only rule.Engine — un-gates MTP and fixes
seq_len:resolve_async_scheduling()no longer excludesPyptoDeepSeekV4Executor; async defaults on for all executors.DecodeRequest.seq_lenderived from the step's own snapshot (sr.num_computed_tokens + sr.num_new_tokens) instead of the placeholder-inflatedreq.num_tokens.Runner — authoritative
seq_lencorrection for MTP (npu_runner.py):_DeepSeekV4MtpRequestStategainscommitted_count(running total of committed output tokens, advances 1-2 per step) andprompt_len(captured at MTP initialization)._main_speculative_batchcorrectsseq_lenstoprompt_len + committed_count + 1for requests that have completed at least one decode (proposed_tokens > 0). The first decode trusts the engine's value. This is needed because under async pipelining the engine must optimistically assume the max tokens per step, so itsseq_lenis inflated whenever a draft is rejected — andtail_position(the runner's previous correction source) stalls on rejection, trapping the model at the same positions. The committed count always advances.2.
refactor(serving): drop dead prev-token plumbingThe serving layer computed a prev-token per decode request and an extra embedding tensor — all discarded by the model layer.
_main_speculative_batchrebuilds bothprev_token_idsandprev_hidden_statesbefore the kernel sees them. Removed from IPC (DecodeRequest.prev_token), engine, worker, and the per-steplookup_embeddings()call.DecodeBatch.prev_token_ids/prev_hidden_statesstay (the MTP runner populates them internally). Net −97 lines.Verification
test_deepseek_v4.pyfailures on this base are unrelated (verified identical on cleanmain; need a rebuiltpypto-lib)