deepinfer-worker
deepinfer/deepinfer/worker/worker.py · a679bb8 · L89–188 · 固定 commit 上游
这是固定版本源码的只读摘录；可在线阅读，也可通过上游链接查看原文件。
89 def execute_model_step(
90 self,
91 is_prompt: bool,
92 is_placeholder: bool,
93 seq_metadata_list: list[SequenceMetadata] | None = None,
94 scheduler_outputs: SchedulerOutputs | None = None,
95 ) -> np.ndarray[np.uint64]:
96 """The unified entry for model execution: every step (prompt/decode, ZOS on/off,
97 placeholder) goes through this Forward-Post_process-Sample round.
98 1. Forward current step (main model)
99 2. Post-process previous step
100 3. Sample current step
101 4. Post-process current step: in place for prompt / non-ZOS steps;
102 ZOS decode steps defer it to the next step to overlap with forward
103
104 Args:
105 is_prompt: Whether this is a prefill step.
106 is_placeholder: Whether this is a placeholder step (no real seqs).
107 seq_metadata_list: Seqs of this step without padding. If is_placeholder, seq_metadata_list is empty.
108 scheduler_outputs: Scheduler outputs for the current step.
109
110 Returns:
111 Request outputs (serialized as uint64 pointers) ready to send back.
112 """
113 # 1. forward curr step
114 total_q = None
115 correction: _SeqlensCorrection | None = None
116 if self.is_driver_worker and not is_placeholder:
117 total_q = sum(seq.seqlen_q for seq in seq_metadata_list) if is_prompt else scheduler_outputs.num_batched_tokens
118 correction = self._correct_seqlens_qk(seq_metadata_list, is_prompt, total_q)
119
120 metadata_zos = correction.metadata_zos if correction is not None else None
121 if not is_placeholder:
122 main_model_output = self.model_runner.execute_main_model(
123 seq_metadata_list, self.gpu_cache, metadata_zos=metadata_zos, total_q=total_q
124 )
125 else:
126 seq_metadata_list = [self.model_runner._placeholder_seq_metadata(is_prompt)]
127 main_model_output = self.model_runner.execute_main_model(seq_metadata_list, self.gpu_cache, metadata_zos=metadata_zos)
128
129 graph_size = main_model_output.graph_size
130 record_runtime = bool(
131 self.is_driver_worker
132 and envs.DEEPINFER_ENABLE_MTP_VARLEN
133 and self.nextn > 0
134 and not is_prompt
135 and not is_placeholder
136 and graph_size > 0
137 )
138 forward_events = cuda_timer.latest_event_pair("device_forward") if record_runtime else None
139 input_metadata = main_model_output.input_metadata
140 model_output = Movable(main_model_output.model_output)
141 del main_model_output
142
143 # 2. post process prev step
144 request_outputs = self._post_process_step(self.prev_step_state)
145
146 # 3. sample curr step（先用 d2h 回来的 seqlens_q 修正 CPU 侧元数据）
147 if correction is not None:
148 if correction.d2h_evt is not None:
149 correction.d2h_evt.synchronize()
150 scheduler_outputs.update_metadata_for_sampling_zos(
151 seq_metadata_list, correction.seqlens_q_cpu, correction.metadata_zos.remap_cpu, self.nextn
152 )
153
154 sampler_output = self.model_runner.execute_sample(
155 seq_metadata_list,
156 input_metadata,
157 model_output,
158 self.gpu_cache,
159 )
160
161 sample_events = cuda_timer.latest_event_pair("device_sample") if record_runtime else None
162 with cpu_timer.record("pre_sampler_sync"):
163 if self.is_driver_worker:
164 self.engine_callbacks.pre_sampler_sync_fn(scheduler_outputs)
165 self.maybe_gc_collect()
166
167 sample_end_evt = torch.cuda.Event(enable_timing=True)
168 sample_end_evt.record()
169
170 curr_step_state = WorkerStepState(
171 seq_metadata_list=seq_metadata_list,
172 scheduler_outputs=scheduler_outputs,
173 sampler_output=sampler_output,
174 is_prompt=is_prompt,
175 is_placeholder=is_placeholder,
176 sample_end_evt=sample_end_evt,
177 forward_events=forward_events,
178 sample_events=sample_events,
179 graph_size=graph_size,
180 )
181
182 # 4. prompt 步或配置 non-ZOS 时当场 post process；ZOS decode 步留到下一步（与 forward 重叠）
183 if is_prompt or not envs.DEEPINFER_ZERO_OVERHEAD_SCHEDULING:
184 request_outputs = np.concatenate([request_outputs, self._post_process_step(curr_step_state)])
185 self.prev_step_state = None
186 else:
187 self.prev_step_state = curr_step_state
188