Overlap the Pipeline
Overlap execution and state processing
Read ZOS forward, postprocess, and sampling order, including one-step-late state correction.
Asynchrony starts with dependencies: launch ready work, then wait only when its result is needed.
01Serial waits leave gaps on the device timeline
If each round waits for samples, updates requests, rebuilds work, and only then launches the next forward, execution waits during control work. Asynchronous scheduling overlaps these stages. The key is identifying available data, state that can be corrected early, and completion dependencies rather than merely adding async to a function.
# Framework pseudocode for ZOS dependencies and deferred state.
previous = None
for plan in scheduler:
metadata = prepare_inputs(plan)
if previous is not None:
metadata = correct_on_device(metadata, previous)
forward = launch_forward(metadata) # preserve execution dependencies
if previous is not None:
wait_for(previous.sample_end)
process_outputs(previous)
sample = launch_sample(forward, metadata)
current = StepState(plan, sample, record_event())
if plan.is_prefill or not zos_enabled:
wait_for(current.sample_end)
process_outputs(current)
previous = None
else:
previous = current
flush_if_present(previous)02Launch this forward before processing the prior result
DeepInfer Worker follows a unified order: current forward, previous postprocess, then current sampling. ZOS decode defers its postprocess to the next round so it overlaps that round’s forward. Prefill and non-ZOS steps process results immediately. At the end, no next round exists to borrow, so an explicit flush is necessary or the final generated results may never reach the request layer.
03Stale lengths cannot be used unchanged
Before previous-step post-processing, control-side sequence lengths can lag, especially while draft-acceptance counts are uncommitted. Upstream _correct_seqlens_qk corrects lengths and inputs from sampling results and constrains synchronization boundaries. Use the pseudocode to read dependencies, then inspect Worker for events, metadata, and ownership.
04Cross-round overlap retains events and state ownership
Official production design: The official deployment article overlaps scheduling, execution, and result processing across rounds. Annotate which round each stage consumes, who owns its results, and which state remains uncommitted. Launching the next forward does not mean the previous result is already visible to the request layer. Public-code check: WorkerStepState construction retains sampling results together with a completion event. _post_process_step waits for that event before consuming results and calling the engine. ZOS decode passes this state to the following round, and flush drains it at the end. Overlap therefore needs readiness signals and an explicit lifetime for in-flight state as well as stage ordering.
Make it yours
Draw three rounds of control and execution with forward, sample, postprocess, and flush. If the prior round accepted one draft, identify metadata that needs correction.