[PLE] The pull op resolves live pointers by layer name

The wait/copy op took host pointers as Python-int arguments, which were baked
into the AOT compile cache as the tracing process's addresses; the first
forward after a cache-load boot then failed. The op now carries only
process-stable values and resolves pointers from a per-process registry.

Co-Authored-By: Hinata <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

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,
                 )