diff --git a/vllm/model_executor/layers/ple_offload_layer.py b/vllm/model_executor/layers/ple_offload_layer.py index d1f668380b..ff581b205a 100644 --- a/vllm/model_executor/layers/ple_offload_layer.py +++ b/vllm/model_executor/layers/ple_offload_layer.py @@ -248,11 +248,46 @@ direct_register_custom_op( # and the consumers read device memory as before. The flag is NOT reset here: # ``release_outputs`` resets it after the forward (unchanged semantics), so # capture warm-ups and dummy runs behave exactly as with ``ple_offload_wait``. -# Pointers are passed as ints so no CPU tensor enters the compiled graph. +# No CPU tensor and no raw address enters the compiled graph: the op resolves +# the live pointers by layer name (see _PULL_LAYERS). # --------------------------------------------------------------------------- +# Fork (found on the reference host): the host flag / host output pointers and +# the row size are NOT op arguments any more. As Python ints they were +# specialised into the compiled graph as constants, so a boot that LOADED the +# AOT compile cache (no Dynamo pass) waited on the previous process's addresses: +# cuStreamWaitValue32 -> CUDA_ERROR_INVALID_VALUE on the first forward. The op +# now carries only process-stable values (buffer, token count, layer name) and +# resolves the live pointers from this registry on every call. +_PULL_LAYERS: dict[str, "PleOffloadLayer"] = {} + + def _ple_offload_wait_copy_impl( + gpu_output_buffer: torch.Tensor, + hidden_states: torch.Tensor, + num_tokens: int, + layer_name: str, +) -> None: + """Resolve this process's live pointers for ``layer_name``, then wait+DMA.""" + layer = _PULL_LAYERS.get(layer_name) + if layer is None or layer._host_output_ptr is None: + raise RuntimeError( + f"PLE wait_copy: no armed pull layer named {layer_name!r} in this " + f"process (armed: {sorted(_PULL_LAYERS)})" + ) + _ple_offload_wait_copy_raw( + gpu_output_buffer, + hidden_states, + layer._sem.flag_ptr, + layer._host_output_ptr, + layer._host_row_bytes, + num_tokens, + layer_name, + ) + + +def _ple_offload_wait_copy_raw( gpu_output_buffer: torch.Tensor, hidden_states: torch.Tensor, flag_ptr: int, @@ -298,9 +333,6 @@ def _ple_offload_wait_copy_impl( def _ple_offload_wait_copy_fake( gpu_output_buffer: torch.Tensor, hidden_states: torch.Tensor, - flag_ptr: int, - host_ptr: int, - row_bytes: int, num_tokens: int, layer_name: str, ) -> None: @@ -395,6 +427,8 @@ class PleOffloadLayer(nn.Module, ABC): self._sem = semaphore self._host_output_ptr = host_output_ptr self._host_row_bytes = int(host_row_bytes) + if host_output_ptr is not None: + _PULL_LAYERS[self._ple_offload_name] = self def forward( self, @@ -409,9 +443,6 @@ class PleOffloadLayer(nn.Module, ABC): torch.ops.vllm.ple_offload_wait_copy( self._gpu_output_buffer, hidden_states, - self._sem.flag_ptr, - self._host_output_ptr, - self._host_row_bytes, input_ids.shape[0], self._ple_offload_name, )