5080deploy: run the PLE graph fill on its own thread, submitted before the launch (fixes #65)

python/freetoken/models/qwen4_exp/ple_disk.py

@@ -6,6 +6,9 @@
 from __future__ import annotations
 
 import os
+import queue
+import threading
+from concurrent.futures import Future
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Sequence
@@ -100,6 +103,52 @@ def resolve_row_source(folder: str) -> PleRowSource:
     return source_from_safetensors(folder)
 
 
+class _Filler:
+    """The one thread that drives the PleStore (engine-thread-only C++, no locks), FIFO.
+
+    Graph fills run here, submitted before the launch: the fill releases the captured decode's
+    WAIT, and cuGraphLaunch may block until the GPU drains, so it must not wait for the launch
+    to return (a large graph deadlocked that way; fork issue #65)."""
+
+    def __init__(self) -> None:
+        self._jobs: queue.SimpleQueue = queue.SimpleQueue()
+        self._thread = threading.Thread(target=self._run, name="ple-filler", daemon=True)
+        self._thread.start()
+
+    def submit(self, fn) -> Future:
+        fut: Future = Future()
+        self._jobs.put((fn, fut))
+        return fut
+
+    def _run(self) -> None:
+        while (item := self._jobs.get()) is not None:
+            fn, fut = item
+            try:
+                fut.set_result(fn())
+            except BaseException as exc:  # noqa: BLE001 -- surfaced to the engine thread
+                fut.set_exception(exc)
+
+    def close(self) -> None:
+        self._jobs.put(None)
+        self._thread.join(timeout=5)
+
+
+class _PendingFill:
+    """A graph fill submitted before its launch; ``cancel`` after a failed launch."""
+
+    def __init__(self, table: "DiskRowTable", future: Future, cancelled: threading.Event) -> None:
+        self._table, self._future, self._cancelled = table, future, cancelled
+
+    def cancel(self) -> None:
+        # no WAIT consumed the signal: clear it, or the next step's WAIT passes on stale rows
+        self._cancelled.set()
+        try:
+            self._future.result()
+        except BaseException:  # noqa: BLE001 -- the launch error is the one to surface
+            pass
+        self._table._flag.zero_()
+
+
 class DiskRowTable:
     """``PLETableBackend`` whose rows are read from disk per fill (--ple-backend disk)."""
 
@@ -185,6 +234,31 @@ def _probe_wait_sync(self, mode: str) -> bool:
 
     # ---------------- host side (engine thread, before the forward launches) ----------------
 
+    def _filler(self) -> _Filler:
+        filler = self.__dict__.get("_filler_thread")
+        if filler is None:
+            filler = self._filler_thread = _Filler()
+        return filler
+
+    def close(self) -> None:
+        filler = self.__dict__.pop("_filler_thread", None)
+        if filler is not None:
+            filler.close()
+
+    def _current_stream(self):
+        return torch.cuda.current_stream(self._device)
+
+    def _raise_failed_fill(self) -> None:
+        pending = self.__dict__.pop("_pending_fill", None)
+        if pending is not None and pending.done() and pending.exception() is not None:
+            raise pending.exception()
+        if pending is not None and not pending.done():
+            self._pending_fill = pending
+
+    def _fill_now(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None:
+        """A fill that completes before its launch, on the store's thread (after any queued graph fill)."""
+        self._filler().submit(lambda: self.fill(runs, graph=graph)).result()
+
     def fill(self, runs: Sequence[torch.Tensor], *, graph: bool) -> None:
         """Stage per-request token runs (two context ids, then the new tokens) in batch order."""
         pinned = self._graph_pinned if graph else self._eager_pinned
@@ -204,22 +278,27 @@ def _ple_context(self, ids: torch.Tensor, position: int) -> list[int]:
             return _context(ids, position, self.eos_token_id)
         return [self.image_token_id if t >= MM_PAD_SHIFT_VALUE else t for t in _context(ids, position, self.eos_token_id)]
 
-    def host_fill_batch(self, batch: Batch, use_graph: bool):
-        """Stage this batch's rows; returns the post-dispatch fill callable under flag-sync, else None."""
-        eos = self.eos_token_id
+    def host_fill_batch(self, batch: Batch, use_graph: bool) -> _PendingFill | None:
+        """Stage this batch's rows before the dispatch. Under flag-sync a graph fill is only submitted:
+        it completes on the filler thread, independently of the launch; returns its handle, else None."""
+        self._raise_failed_fill()
         if batch.is_decode:
             reqs = list(batch.reqs)
+            # the context comes from the request as it is NOW: the engine mutates it once the launch returns
+            contexts = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs]
             if use_graph and self._wait_sync:
                 bs = batch.padded_size
                 self._token_readback[:bs].copy_(batch.input_ids, non_blocking=True)
-                self._readback_event.record(torch.cuda.current_stream(self._device))
+                self._readback_event.record(self._current_stream())
+                cancelled = threading.Event()
 
-                def _complete() -> None:
+                def _job() -> None:
                     try:
                         self._readback_event.synchronize()
                         tokens = self._token_readback[:bs].to(torch.int64).tolist()
-                        runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64)
-                                for r, t in zip(reqs, tokens)]
+                        if cancelled.is_set():
+                            return
+                        runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)]
                         self.fill(runs, graph=True)
                     except BaseException:
                         from freetoken.kernel import _ple_store
@@ -228,12 +307,12 @@ def _complete() -> None:
                         _ple_store.signal_flag(self._flag.data_ptr())
                         raise
 
-                return _complete
+                self._pending_fill = self._filler().submit(_job)
+                return _PendingFill(self, self._pending_fill, cancelled)
             # launch-gating: this D2H is the step's readback and orders the fill after sampling
             tokens = batch.input_ids.to("cpu").to(torch.int64).tolist()
-            runs = [torch.tensor([*self._ple_context(r.input_ids, r.device_len - 1), t], dtype=torch.int64)
-                    for r, t in zip(reqs, tokens)]
-            self.fill(runs, graph=use_graph)
+            runs = [torch.tensor([*ctx, t], dtype=torch.int64) for ctx, t in zip(contexts, tokens)]
+            self._fill_now(runs, graph=use_graph)
             return None
         runs = [
             torch.cat((
@@ -242,17 +321,19 @@ def _complete() -> None:
             ))
             for req in batch.padded_reqs
         ]
-        self.fill(runs, graph=False)
+        self._fill_now(runs, graph=False)
         return None
 
     @contextmanager
     def forward_host_ctx(self, batch: Batch, use_graph: bool):
-        """Around one dispatch: stage on enter, run the deferred fill+signal on exit."""
-        deferred = self.host_fill_batch(batch, use_graph)
-        yield
-        # no try/finally: a failed launch leaves no WAIT pending, so the fill must not run
-        if deferred is not None:
-            deferred()
+        """Around one dispatch: stage (or submit the graph fill) on enter; cancel a submitted fill if the launch fails."""
+        pending = self.host_fill_batch(batch, use_graph)
+        try:
+            yield
+        except BaseException:
+            if pending is not None:
+                pending.cancel()
+            raise
 
     # ---------------- device side (PLETableBackend protocol) ----------------
 

tests/models/qwen4_exp/test_ple_disk_fill_order.py

@@ -0,0 +1,134 @@
+"""The PLE disk backend's wait-sync fill must not depend on the graph launch returning.
+
+The captured decode graph WAITs on ``_flag``; ``cuGraphLaunch`` may block until the GPU drains
+(a large graph fills the launch queue), and the GPU is parked on that WAIT. So the fill that sets
+the flag has to run on its own thread, started before the launch (fork issue #65).
+These tests stand in for the GPU with host fakes: the "launch" blocks until the flag is set.
+"""
+
+from __future__ import annotations
+
+import threading
+import time
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+import freetoken.models.qwen4_exp.ple_disk as ple_disk
+from freetoken.models.qwen4_exp.ple_disk import DiskRowTable
+
+EOS = 9
+
+
+class _Store:
+    """Stands in for the C++ PleStore: records calls and the thread that made them."""
+
+    def __init__(self, flag: torch.Tensor):
+        self.flag = flag
+        self.calls: list[tuple[str, list[int] | None, str]] = []
+
+    def stage(self, tokens_addr, n, staging_addr):
+        tokens = ctypes_read(tokens_addr, n + 2)
+        self.calls.append(("stage", tokens, threading.current_thread().name))
+
+    def flush(self, signal_addr):
+        self.calls.append(("flush", None, threading.current_thread().name))
+        if signal_addr:
+            self.flag[0] = 1
+
+
+def ctypes_read(addr: int, n: int) -> list[int]:
+    import ctypes
+
+    return list((ctypes.c_int64 * n).from_address(addr))
+
+
+class _Event:
+    def record(self, stream=None):
+        pass
+
+    def synchronize(self):
+        pass
+
+
+def _table(monkeypatch) -> DiskRowTable:
+    # host-only: no CUDA stream behind the (faked) readback event
+    monkeypatch.setattr(ple_disk.torch.cuda, "current_stream", lambda *a, **k: None)
+    t = DiskRowTable.__new__(DiskRowTable)
+    t._device = None
+    t._wait_sync = True
+    t._flag = torch.zeros(1, dtype=torch.int64)
+    t._store = _Store(t._flag)
+    t._token_readback = torch.zeros(8, dtype=torch.int32)
+    t._readback_event = _Event()
+    t._graph_pinned = torch.zeros(64, dtype=torch.uint8)
+    t._eager_pinned = torch.zeros(64, dtype=torch.uint8)
+    t._token_bytes = 1
+    t.eos_token_id = EOS
+    t.image_token_id = None
+    return t
+
+
+def _close(t: DiskRowTable) -> None:
+    close = getattr(t, "close", None)
+    if close is not None:
+        close()
+
+
+def _decode_batch(prev_tokens: list[int], new_token: int):
+    req = SimpleNamespace(input_ids=torch.tensor(prev_tokens + [new_token]), device_len=len(prev_tokens) + 1)
+    return SimpleNamespace(is_decode=True, reqs=[req], padded_size=1,
+                           input_ids=torch.tensor([new_token], dtype=torch.int32))
+
+
+def _blocking_launch(flag: torch.Tensor, timeout: float = 5.0) -> bool:
+    """A cuGraphLaunch that cannot return until the GPU passes the graph's WAIT(flag)."""
+    deadline = time.monotonic() + timeout
+    while time.monotonic() < deadline:
+        if int(flag[0]) == 1:
+            flag[0] = 0  # the graph's WAIT-then-RESET
+            return True
+        time.sleep(0.001)
+    return False
+
+
+def test_graph_fill_does_not_wait_for_the_launch_to_return(monkeypatch):
+    t = _table(monkeypatch)
+    with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True):
+        assert _blocking_launch(t._flag), "deadlock: the fill only runs after the launch returns"
+    _close(t)
+
+
+def test_fills_run_in_step_order_on_one_thread_with_contexts_taken_before_the_launch(monkeypatch):
+    t = _table(monkeypatch)
+    for prev, new in (([1, 2, 3], 7), ([1, 2, 3, 7], 8), ([1, 2, 3, 7, 8], 4)):
+        batch = _decode_batch(prev, new)
+        with t.forward_host_ctx(batch, use_graph=True):
+            # the engine mutates the request as soon as the launch returns (overlap scheduling)
+            batch.reqs[0].input_ids = torch.tensor([0, 0, 0, 0, 0, 0])
+            batch.reqs[0].device_len += 1
+            assert _blocking_launch(t._flag)
+    _close(t)
+    staged = [c for c in t._store.calls if c[0] == "stage"]
+    assert [c[1] for c in staged] == [[2, 3, 7], [3, 7, 8], [7, 8, 4]]
+    assert len({c[2] for c in t._store.calls}) == 1  # the store stays single-threaded
+
+
+def test_eager_fill_completes_before_the_launch(monkeypatch):
+    t = _table(monkeypatch)
+    req = SimpleNamespace(input_ids=torch.tensor([1, 2, 3, 4]), cached_len=1, device_len=4)
+    batch = SimpleNamespace(is_decode=False, padded_reqs=[req])
+    with t.forward_host_ctx(batch, use_graph=False):
+        assert [c[0] for c in t._store.calls] == ["stage", "flush"]
+    _close(t)
+
+
+def test_a_failed_launch_leaves_no_stale_signal(monkeypatch):
+    t = _table(monkeypatch)
+    with pytest.raises(RuntimeError, match="launch failed"):
+        with t.forward_host_ctx(_decode_batch([1, 2, 3], 7), use_graph=True):
+            raise RuntimeError("launch failed")
+    # no WAIT consumed it, so a leftover 1 would let the next step read stale rows
+    assert int(t._flag[0]) == 0
+    _close(t)