From 6ebfa18e2071714900dbbb7705aafbd03c10edf4 Mon Sep 17 00:00:00 2001 From: "Gabriel A. Devenyi" Date: Wed, 23 Sep 2026 23:20:41 -0400 Subject: [PATCH] fix(qwen4_exp): run the PLE graph fill on its own thread, submitted before the launch The disk PLE table's wait-sync mode captures a WAIT on _flag in every decode graph and set the flag only after the graph launch returned. cuGraphLaunch can block until the GPU drains when the graph does not fit the launch queue -- the GPU parked on that WAIT -- so a large decode graph (hybrid MoE; hybrid + --moe-collect-stats every time) deadlocked on its first step (fork #65). All PleStore calls now run FIFO on one filler thread (the store stays single-threaded). A graph fill is submitted before the launch with the request contexts snapshotted on the engine thread; eager fills are submitted and awaited as before. A failed launch cancels its fill and clears a signal no WAIT consumed. Assisted-by: Claude Opus 5.5 --- python/freetoken/models/qwen4_exp/ple_disk.py | 117 ++++++++++++--- .../qwen4_exp/test_ple_disk_fill_order.py | 134 ++++++++++++++++++ 2 files changed, 233 insertions(+), 18 deletions(-) create mode 100644 tests/models/qwen4_exp/test_ple_disk_fill_order.py diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index cea6269357..223b97ef66 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/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) ---------------- diff --git a/tests/models/qwen4_exp/test_ple_disk_fill_order.py b/tests/models/qwen4_exp/test_ple_disk_fill_order.py new file mode 100644 index 0000000000..42dbee9949 --- /dev/null +++ b/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)