fix(ple): add rocm disk wait-sync support

python/freetoken/kernel/csrc/row_store/row_store_ext.cpp

@@ -8,10 +8,11 @@
 //                Hash reference: tests/models/qwen4_exp/test_ple_disk.py.
 //
 // Platform seams: TableFile (O_DIRECT+pread; Win: NO_BUFFERING), BatchReader (io_uring,
-// pread-pool fallback = the portable shape), cumemop_* (dlopen libcuda; Win: nvcuda).
+// pread-pool fallback = the portable shape), stream memops (CUDA driver or HIP signal memory).
 
 #include <algorithm>
 #include <cerrno>
+#include <chrono>
 #include <condition_variable>
 #include <cstdint>
 #include <cstdlib>
@@ -41,6 +42,10 @@
 
 #include <torch/extension.h>
 
+#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM)
+#include <hip/hip_runtime.h>
+#endif
+
 namespace py = pybind11;
 
 namespace {
@@ -63,6 +68,109 @@ uint8_t *page_aligned_alloc(size_t bytes) {
   return static_cast<uint8_t *>(p);
 }
 
+#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM)
+
+hipError_t alloc_signal(int64_t **out) {
+  void *word = nullptr;
+  // Like CPU MoE, HIP stream waits require signal memory, not a pinned staging allocation.
+  const auto rc = hipExtMallocWithFlags(&word, sizeof(int64_t), hipMallocSignalMemory);
+  if (rc == hipSuccess) {
+    *out = static_cast<int64_t *>(word);
+    release_store_i64(*out, 0);
+  }
+  return rc;
+}
+
+torch::Tensor alloc_signal_tensor() {
+  int64_t *word = nullptr;
+  const auto rc = alloc_signal(&word);
+  TORCH_CHECK(rc == hipSuccess, "HIP signal allocation failed: ", hipGetErrorString(rc));
+  return torch::from_blob(word, {1}, [](void *p) { (void)hipFree(p); },
+                          torch::TensorOptions().dtype(torch::kInt64).device(torch::kCPU));
+}
+
+int memop_write(uintptr_t stream, uintptr_t addr, uint64_t value) {
+  return hipStreamWriteValue64(reinterpret_cast<hipStream_t>(stream),
+                               reinterpret_cast<void *>(addr), value, hipStreamWriteValueDefault);
+}
+
+int memop_wait_geq(uintptr_t stream, uintptr_t addr, uint64_t value) {
+  return hipStreamWaitValue64(reinterpret_cast<hipStream_t>(stream),
+                              reinterpret_cast<void *>(addr), value, hipStreamWaitValueGte, UINT64_MAX);
+}
+
+// Capability failures must not poison gating fallback or hide a caller's earlier error.
+struct ProbeErrorScope {
+  ProbeErrorScope() {
+    TORCH_CHECK(hipPeekAtLastError() == hipSuccess, "HIP error pending before row-store sync probe");
+  }
+  ~ProbeErrorScope() { (void)hipGetLastError(); }
+};
+
+bool capture_memops_supported() {
+  ProbeErrorScope errors;
+  int device = 0;
+  int can_wait = 0;
+  if (hipGetDevice(&device) != hipSuccess ||
+      hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != hipSuccess || can_wait == 0)
+    return false;
+
+  int64_t *ready = nullptr;
+  int64_t *done = nullptr;
+  hipStream_t stream = nullptr;
+  hipGraph_t graph = nullptr;
+  hipGraphExec_t exec = nullptr;
+  auto rc = alloc_signal(&ready);
+  if (rc == hipSuccess) rc = alloc_signal(&done);
+  if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking);
+  // A runtime that executes the WAIT during capture must not block the probe.
+  if (rc == hipSuccess) release_store_i64(done, 1);
+  if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed);
+  const bool started = rc == hipSuccess;
+  if (rc == hipSuccess) rc = static_cast<hipError_t>(memop_write(
+      reinterpret_cast<uintptr_t>(stream), reinterpret_cast<uintptr_t>(ready), 1));
+  if (rc == hipSuccess) rc = static_cast<hipError_t>(memop_wait_geq(
+      reinterpret_cast<uintptr_t>(stream), reinterpret_cast<uintptr_t>(done), 1));
+  if (started) {
+    const auto end_rc = hipStreamEndCapture(stream, &graph);
+    if (rc == hipSuccess) rc = end_rc;
+  }
+  if (rc == hipSuccess) rc = hipStreamSynchronize(stream);
+  size_t nodes = 0;
+  if (rc == hipSuccess) rc = hipGraphGetNodes(graph, nullptr, &nodes);
+  // Older HIP runtimes execute memops during capture without recording them.
+  if (rc == hipSuccess && nodes != 2) rc = hipErrorNotSupported;
+  if (rc == hipSuccess) rc = hipGraphInstantiate(&exec, graph, nullptr, nullptr, 0);
+  bool served = false;
+  if (rc == hipSuccess) {
+    release_store_i64(ready, 0);
+    release_store_i64(done, 0);
+    std::thread responder([&] {
+      const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2);
+      while (__atomic_load_n(ready, __ATOMIC_ACQUIRE) == 0 && std::chrono::steady_clock::now() < deadline)
+        std::this_thread::yield();
+      served = __atomic_load_n(ready, __ATOMIC_ACQUIRE) == 1;
+      // The graph never resets done, so even a launch arriving after the deadline is released.
+      release_store_i64(done, 1);
+    });
+    rc = hipGraphLaunch(exec, stream);
+    if (rc == hipSuccess) rc = hipStreamSynchronize(stream);
+    responder.join();
+  }
+  const bool ok = rc == hipSuccess && served;
+  if (exec != nullptr) (void)hipGraphExecDestroy(exec);
+  if (graph != nullptr) (void)hipGraphDestroy(graph);
+  if (stream != nullptr) {
+    (void)hipStreamSynchronize(stream);
+    (void)hipStreamDestroy(stream);
+  }
+  if (ready != nullptr) (void)hipFree(ready);
+  if (done != nullptr) (void)hipFree(done);
+  return ok;
+}
+
+#else
+
 // Stream memops for the flag-sync fast path, resolved from the driver at runtime.
 using CuMemOp64Fn = int (*)(void *stream, unsigned long long addr, unsigned long long value,
                             unsigned int flags);
@@ -97,6 +205,8 @@ int memop_wait_geq(uintptr_t stream, uintptr_t addr, uint64_t value) {
   return g_cu_wait64(reinterpret_cast<void *>(stream), addr, value, kCuWaitValueGeq);
 }
 
+#endif
+
 // WAIT(>=1) then RESET: resetting first would race a fast host signal and deadlock the stream.
 void memop_wait_reset(uintptr_t stream, uintptr_t flag_addr) {
   if (memop_wait_geq(stream, flag_addr, 1) != 0 || memop_write(stream, flag_addr, 0) != 0)
@@ -674,4 +784,8 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("memop_wait_geq", &memop_wait_geq, py::arg("stream"), py::arg("addr"), py::arg("value"));
   m.def("memop_wait_reset", &memop_wait_reset, py::arg("stream"), py::arg("flag_addr"));
   m.def("signal_flag", &signal_flag, py::arg("flag_addr"));
+#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM)
+  m.def("alloc_signal_tensor", &alloc_signal_tensor);
+  m.def("capture_memops_supported", &capture_memops_supported);
+#endif
 }

python/freetoken/kernel/row_store.py

@@ -3,9 +3,10 @@
 ``RowStore`` reads fixed-width rows straight out of checkpoint shard files (io_uring or a pread pool,
 O_DIRECT) into pinned staging by row id; ``PleStore`` adds PLE's n-gram hashing.
 
-Graph sync: a captured graph cannot wait for host I/O, so a captured ``lookup`` WAITs on a pinned
-flag through a stream memop (``wait_reset``) that the host ``signal``s once the fill landed; where
-the driver rejects memops in capture, the caller falls back to launch-gating (fill before replay).
+Graph sync: a captured graph cannot wait for host I/O, so a captured ``lookup`` WAITs on a host-visible
+flag (pinned on CUDA, signal memory on HIP) through a stream memop (``wait_reset``). The host
+``signal``s once the fill landed; where the driver rejects memops in capture, the caller falls back
+to launch-gating (fill before replay).
 ``probe_wait_sync`` decides which, once, at startup.
 """
 
@@ -20,22 +21,32 @@
 PleStore = _row_store.PleStore
 
 
+def alloc_sync_flag() -> torch.Tensor:
+    if torch.version.hip is not None:
+        return _row_store.alloc_signal_tensor()
+    return alloc_pinned_tensor(1, dtype=torch.int64).zero_()
+
+
 def probe_wait_sync(mode: str, device: torch.device) -> bool:
     """``mode``: ``auto`` (probe), ``wait`` (require memops) or ``gate`` (never use them)."""
     if mode not in ("auto", "wait", "gate"):
         raise ValueError(f"unknown row-store sync mode {mode!r}; expected auto, wait or gate")
     if mode == "gate":
         return False
-    scratch = alloc_pinned_tensor(1, dtype=torch.int64)
-    scratch.zero_()
-    stream = torch.cuda.current_stream(device)
-    ok = (
-        _row_store.memop_write(stream.cuda_stream, scratch.data_ptr(), 7) == 0
-        and _row_store.memop_wait_geq(stream.cuda_stream, scratch.data_ptr(), 7) == 0
-    )
-    if ok:
-        stream.synchronize()
-        ok = int(scratch[0]) == 7
+    if torch.version.hip is not None:
+        with torch.cuda.device(device):
+            ok = _row_store.capture_memops_supported()
+    else:
+        scratch = alloc_pinned_tensor(1, dtype=torch.int64)
+        scratch.zero_()
+        stream = torch.cuda.current_stream(device)
+        ok = (
+            _row_store.memop_write(stream.cuda_stream, scratch.data_ptr(), 7) == 0
+            and _row_store.memop_wait_geq(stream.cuda_stream, scratch.data_ptr(), 7) == 0
+        )
+        if ok:
+            stream.synchronize()
+            ok = int(scratch[0]) == 7
     if mode == "wait" and not ok:
         raise RuntimeError("wait-sync requested but stream memops are unavailable")
     return ok
@@ -51,4 +62,4 @@ def signal(flag: torch.Tensor) -> None:
     _row_store.signal_flag(flag.data_ptr())
 
 
-__all__ = ["RowStore", "PleStore", "probe_wait_sync", "wait_reset", "signal"]
+__all__ = ["RowStore", "PleStore", "alloc_sync_flag", "probe_wait_sync", "wait_reset", "signal"]

python/freetoken/models/deepseek_v41/engram_table.py

@@ -130,8 +130,7 @@ def __init__(
         self._eager_slot = 0
         self._eager_dev = None
         self._grow_eager(max_extend_tokens)
-        self.flag = alloc_pinned_tensor(1, dtype=torch.int64)
-        self.flag.zero_()
+        self.flag = None
         self.wait_sync = False  # set by the host coordinator after its probe
 
     def _grow_eager(self, num_tokens: int) -> None:
@@ -209,6 +208,11 @@ def __init__(self, hash: EngramHash, tables: Sequence[EngramDiskTable], device:
         self.wait_sync = probe_wait_sync(sync_mode, device)
         for t in self.tables:
             t.wait_sync = self.wait_sync
+            if self.wait_sync:
+                from freetoken.kernel.row_store import alloc_sync_flag
+
+                with torch.cuda.device(device):
+                    t.flag = alloc_sync_flag()
         self._token_readback = alloc_pinned_tensor(min(t.max_graph_rows for t in self.tables), dtype=torch.int32)
         self._readback_event = torch.cuda.Event()
         self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="engram-fill")

python/freetoken/models/qwen4_exp/ple_disk.py

@@ -6,6 +6,7 @@
 from __future__ import annotations
 
 import os
+from concurrent.futures import Future, ThreadPoolExecutor
 from contextlib import contextmanager
 from dataclasses import dataclass
 from typing import Sequence
@@ -202,12 +203,58 @@ def __init__(
 
         self._wait_sync = probe_wait_sync(os.getenv(_SYNC_ENV, "auto"), self._device)
         # one flag for all graphs: the readback event orders a fill after the previous graph, so signals never overlap
-        self._flag = alloc_pinned_tensor(1, dtype=torch.int64)
-        self._flag.zero_()
+        from freetoken.kernel.row_store import alloc_sync_flag
+
+        self._flag = alloc_sync_flag() if self._wait_sync else None
         self._token_readback = alloc_pinned_tensor(max_graph_rows, dtype=torch.int32)
         self._readback_event = torch.cuda.Event()
         sync = "wait-sync" if self._wait_sync else "launch-gating"
         logger.info_rank0(f"PLE disk backend: {self._store.io_backend()}, {sync}")
+        self._worker = (
+            ThreadPoolExecutor(max_workers=1, thread_name_prefix="ple-fill")
+            if self._wait_sync and torch.version.hip is not None else None
+        )
+        self._pending: Future | None = None
+        self._failed = False
+
+    def _fill_decode(self, histories: Sequence[Sequence[int]], bs: int) -> None:
+        try:
+            self._readback_event.synchronize()
+            tokens = self._token_readback[:bs].to(torch.int64).tolist()
+            runs = [torch.tensor([*history, t], dtype=torch.int64)
+                    for history, t in zip(histories, tokens)]
+            self.fill(runs, graph=True)
+        except BaseException:
+            from freetoken.kernel.row_store import signal
+
+            # unblock the stream before surfacing; the step's output is discarded
+            signal(self._flag)
+            raise
+
+    def _start_fill(self, histories: Sequence[Sequence[int]], bs: int) -> None:
+        @torch.inference_mode()
+        def _complete() -> None:
+            self._fill_decode(histories, bs)
+
+        # Like Engram, release HIP's WAIT independently of a queue-blocked graph launch.
+        self._pending = self._worker.submit(_complete)
+
+    def _await_pending(self) -> None:
+        if self._pending is not None:
+            pending, self._pending = self._pending, None
+            try:
+                pending.result()
+            except BaseException:
+                self._failed = True
+                if self._worker is not None:
+                    # An interrupted waiter must drain the fill before recording its consumed event.
+                    self._worker.shutdown(wait=True)
+                raise
+
+    def __del__(self) -> None:
+        worker = getattr(self, "_worker", None)
+        if worker is not None:
+            worker.shutdown(wait=False)
 
     # ---------------- host side (engine thread, before the forward launches) ----------------
 
@@ -237,28 +284,25 @@ def _ple_context(self, ids: torch.Tensor, position: int) -> list[int]:
         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."""
+        """Stage this batch's rows; returns the deferred fill for wait-sync without a worker."""
+        if torch.version.hip is not None:
+            self._await_pending()
+            if self._failed:
+                raise RuntimeError("PLE disk table failed; recreate the engine before dispatch")
         eos = self.eos_token_id
         if batch.is_decode:
             reqs = list(batch.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))
+                histories = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs]
+                if self._worker is not None:
+                    self._start_fill(histories, bs)
+                    return None
 
                 def _complete() -> 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)]
-                        self.fill(runs, graph=True)
-                    except BaseException:
-                        from freetoken.kernel.row_store import signal
-
-                        # unblock the stream before surfacing; the step's output is discarded
-                        signal(self._flag)
-                        raise
+                    self._fill_decode(histories, bs)
 
                 return _complete
             # launch-gating: this D2H is the step's readback and orders the fill after sampling
@@ -279,8 +323,22 @@ def _complete() -> 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."""
+        """Around one dispatch: stage on enter, finish the fill+signal on exit."""
         deferred = self.host_fill_batch(batch, use_graph)
+        if self._worker is not None:
+            try:
+                yield
+            except BaseException:
+                self._failed = True
+                raise
+            finally:
+                try:
+                    self._await_pending()
+                finally:
+                    # The worker must finish before recording the event it waits on at the next fill.
+                    if use_graph:
+                        self._graph_consumed.record(torch.cuda.current_stream(self._device))
+            return
         yield
         # no try/finally: a failed launch leaves no WAIT pending, so the fill must not run
         if deferred is not None:

setup.py

@@ -287,7 +287,11 @@ def build_extension(self, ext):
                 sources=[
                     "python/freetoken/kernel/csrc/row_store/row_store_ext.cpp",
                 ],
+                include_dirs=runtime_include_dirs if IS_ROCM else [],
+                library_dirs=runtime_library_dirs if IS_ROCM else [],
+                libraries=[runtime_lib] if IS_ROCM else [],
                 extra_compile_args=["-O3", "-std=c++20"],
+                extra_link_args=runtime_link_args if IS_ROCM else [],
             )
         ] if sys.platform == "linux" else []),
         *ggml_cpu_modules,

tests/models/deepseek_v41/test_engram.py

@@ -101,6 +101,10 @@ def oracle(rows: torch.Tensor) -> torch.Tensor:
     from freetoken.kernel.row_store import probe_wait_sync
 
     table.wait_sync = probe_wait_sync("auto", torch.device("cuda"))
+    if table.wait_sync:
+        from freetoken.kernel.row_store import alloc_sync_flag
+
+        table.flag = alloc_sync_flag()
     dec_rows = hash.row_ids(ids[40 - 1 :], context=ids[: 40 - 1])[:, 0]
     out = torch.empty(1, hash.n_cols * 32, dtype=torch.bfloat16, device="cuda")
     graph = torch.cuda.CUDAGraph()