| |
| |
|
|
| """Eager prefill chunk-sequence orchestration and ownership.""" |
|
|
| from __future__ import annotations |
|
|
| from typing import Any, Callable |
|
|
| from models.common.llm_runtime.prefill.inputs import PrefillDeviceInputs, PrefillInputStager, PrefillPositionInputs |
| from models.common.llm_runtime.prefill.plan import PrefillChunk, PrefillRequest |
| from models.common.llm_runtime.prefill.postprocess import PrefillPostprocessor, retain_owned |
| from models.common.llm_runtime.prefill.result_collector import InvocationResult |
| from models.common.llm_runtime.prefill.signatures import PreparedPrefill |
| from models.common.llm_runtime.tensor_resources import attach_cleanup_failures, raise_cleanup_failures |
|
|
|
|
| class PrefillSequenceRunner: |
| """Execute one prepared eager request while preserving tensor ownership.""" |
|
|
| def __init__( |
| self, |
| *, |
| input_stager: PrefillInputStager, |
| postprocessor: PrefillPostprocessor, |
| run_hidden_body: Callable[[PrefillRequest, PrefillDeviceInputs], Any], |
| run_chunk_body: Callable[[PreparedPrefill, PrefillChunk, PrefillDeviceInputs, PrefillPositionInputs], Any], |
| release_transient: Callable[[Any], list[BaseException]], |
| ) -> None: |
| self.input_stager = input_stager |
| self.postprocessor = postprocessor |
| self.run_hidden_body = run_hidden_body |
| self.run_chunk_body = run_chunk_body |
| self.release_transient = release_transient |
|
|
| def run(self, prepared: PreparedPrefill, *, count_tokens: bool = True) -> InvocationResult: |
| """Execute the request's planned chunks as one eager prefill sequence.""" |
|
|
| request = prepared.request |
| final_chunk = request.chunks[-1] |
| if request.uses_chunked_prefill: |
| final_relative_last = (request.last_token_indices[0] - final_chunk.chunk_start_idx) % final_chunk.chunk_size |
| else: |
| final_relative_last = max( |
| last - cached for last, cached in zip(request.last_token_indices, request.cached_tokens) |
| ) |
|
|
| owned: list[Any] = [] |
| kpt = None |
| kpt_prepared = False |
| final_step_output = None |
| final_position_inputs = None |
| sampled_output = None |
| try: |
| if request.uses_chunked_prefill: |
| kpt = self.postprocessor.make_device_kpt( |
| self.postprocessor.prepared_sampling(prepared), |
| self.postprocessor.sampling_output_rows(prepared), |
| force_topk=prepared.sampling_path == "topk", |
| ) |
| kpt_prepared = True |
| retain_owned(owned, kpt) |
|
|
| for chunk in request.chunks: |
| device_inputs, position_inputs = self.input_stager.stage_step( |
| request, |
| chunk, |
| final_relative_last, |
| ) |
| retain_owned(owned, device_inputs) |
| retain_owned(owned, position_inputs) |
| if not kpt_prepared: |
| kpt = self.postprocessor.make_device_kpt( |
| self.postprocessor.prepared_sampling(prepared), |
| self.postprocessor.sampling_output_rows(prepared), |
| force_topk=prepared.sampling_path == "topk", |
| ) |
| kpt_prepared = True |
| retain_owned(owned, kpt) |
| step_output = self._execute_step( |
| prepared, |
| chunk, |
| device_inputs, |
| position_inputs, |
| ) |
| if chunk.contains_last_token: |
| final_step_output = step_output |
| final_position_inputs = position_inputs |
| retain_owned(owned, final_step_output) |
| break |
| intermediate_output = step_output |
| step_output = None |
| failures = self.release_transient(intermediate_output) |
| if failures: |
| raise_cleanup_failures(failures) |
|
|
| if final_step_output is None or final_position_inputs is None: |
| raise RuntimeError("planned prefill sequence did not produce a final output") |
| if not request.uses_chunked_prefill and prepared.sampling_path == "topk": |
| sampled_output = self.postprocessor.make_sampling_output( |
| self.postprocessor.sampling_output_rows(prepared) |
| ) |
| retain_owned(owned, sampled_output) |
| output = self.postprocessor.finish_prefill_sequence( |
| prepared, |
| final_step_output, |
| kpt, |
| final_position_inputs, |
| sampled_output=sampled_output, |
| owned=owned, |
| count_tokens=count_tokens, |
| ) |
| except BaseException as primary: |
| failures = self.release_transient(tuple(owned)) |
| attach_cleanup_failures(primary, failures) |
| raise |
| return InvocationResult(value=output, owned=tuple(owned)) |
|
|
| def _execute_step( |
| self, |
| prepared: PreparedPrefill, |
| chunk: PrefillChunk, |
| device_inputs: PrefillDeviceInputs, |
| position_inputs: PrefillPositionInputs, |
| ) -> Any: |
| if not prepared.request.uses_chunked_prefill: |
| return self.run_hidden_body(prepared.request, device_inputs) |
| return self.run_chunk_body(prepared, chunk, device_inputs, position_inputs) |
|
|
