s02Inference core22 min read

Inside an Execution Step

Inspect an execution step

Understand how Worker and ModelRunner pass input metadata, cache references, and sampling results.

Framework boundaries follow inputs, state, and dependencies, not just call order.

01Read input metadata before entering the model

The scheduler selects sequences, positions processed per sequence, and cache-block references. ModelRunner organizes them as execution inputs. Read sequence lengths, positions, and cache references before following execute_main_model into the forward entry. The framework must preserve each request’s semantics within shared work.

s02_execution_step.py
# Framework pseudocode summarizing Worker responsibilities.
metadata = prepare_inputs(scheduled_sequences, cache_state)
main_output = model_runner.execute_main_model(metadata)

post_process(previous_step_state)
sampled = model_runner.execute_sample(main_output, metadata)
current = StepState(metadata, sampled)

if is_prompt_step or not zos_enabled:
    ready_outputs = post_process(current)
    previous_step_state = None
else:
    previous_step_state = current

# Concrete metadata and method signatures live upstream.

02Forward and sampling are separate boundaries

execute_main_model returns model output and metadata; execute_sample prepares sampler inputs and calls the model’s sample path. Separating these boundaries accommodates constrained output, draft verification, and state post-processing. The excerpt focuses on who creates, passes, and consumes these objects rather than model-layer mathematics.

03Worker orders the dependencies

Worker.execute_model_step orders current forward, previous-step post-processing, and current sampling. Prefill and non-ZOS paths post-process immediately; ZOS decode retains the current StepState for the following round. Arrows express data dependencies, not a guarantee that current sampling updates request state before the call returns.

TRY IT YOURSELF

Make it yours

Draw two Worker rounds and mark where main_output, sampler_output, and StepState are created and consumed.

Add this mechanism to your toolkit.