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

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

@@ -8,7 +8,7 @@
 //                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>
@@ -41,6 +41,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 +67,70 @@ uint8_t *page_aligned_alloc(size_t bytes) {
   return static_cast<uint8_t *>(p);
 }
 
+#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM)
+
+torch::Tensor alloc_signal_tensor() {
+  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);
+  TORCH_CHECK(rc == hipSuccess, "HIP signal allocation failed: ", hipGetErrorString(rc));
+  release_store_i64(static_cast<int64_t *>(word), 0);
+  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);
+}
+
+bool capture_memops_supported() {
+  TORCH_CHECK(hipPeekAtLastError() == hipSuccess, "HIP error pending before row-store sync probe");
+  auto flag = alloc_signal_tensor();
+  const auto addr = reinterpret_cast<uintptr_t>(flag.data_ptr<int64_t>());
+  release_store_i64(flag.data_ptr<int64_t>(), 7);
+  hipStream_t stream = nullptr;
+  hipGraph_t graph = nullptr;
+  hipGraphExec_t exec = nullptr;
+  auto rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking);
+  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), addr, 7));
+  if (rc == hipSuccess) rc = static_cast<hipError_t>(memop_wait_geq(reinterpret_cast<uintptr_t>(stream), addr, 7));
+  if (started) {
+    const auto end_rc = hipStreamEndCapture(stream, &graph);
+    if (rc == hipSuccess) rc = end_rc;
+  }
+  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);
+  if (rc == hipSuccess) {
+    release_store_i64(flag.data_ptr<int64_t>(), 0);
+    rc = hipGraphLaunch(exec, stream);
+  }
+  if (rc == hipSuccess) rc = hipStreamSynchronize(stream);
+  const bool ok = rc == hipSuccess && *flag.data_ptr<int64_t>() == 7;
+  if (exec != nullptr) (void)hipGraphExecDestroy(exec);
+  if (graph != nullptr) (void)hipGraphDestroy(graph);
+  if (stream != nullptr) {
+    (void)hipStreamSynchronize(stream);
+    (void)hipStreamDestroy(stream);
+  }
+  (void)hipGetLastError();
+  TORCH_CHECK(rc == hipSuccess || rc == hipErrorNotSupported,
+              "HIP row-store sync probe failed: ", hipGetErrorString(rc));
+  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 +165,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 +744,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

@@ -20,22 +20,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 +61,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

@@ -209,6 +209,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 and torch.version.hip is not None:
+                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,71 @@ 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_()
+        if torch.version.hip is not None:
+            from freetoken.kernel.row_store import alloc_sync_flag
+
+            self._flag = alloc_sync_flag() if self._wait_sync else None
+        else:
+            self._flag = alloc_pinned_tensor(1, dtype=torch.int64)
+            self._flag.zero_()
         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
+        self._closed = False
+
+    def _start_fill(self, reqs, bs: int):
+        histories = [self._ple_context(r.input_ids, r.device_len - 1) for r in reqs]
+
+        @torch.inference_mode()
+        def _complete() -> 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
+
+                signal(self._flag)
+                raise
+
+        # Like Engram, release HIP's WAIT independently of a queue-blocked graph launch.
+        self._pending = self._worker.submit(_complete)
+        return self._await_pending
+
+    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 close(self) -> None:
+        self._closed = True
+        if self._worker is not None:
+            self._worker.shutdown(wait=True)
+            try:
+                self._await_pending()
+            finally:
+                self._graph_consumed.synchronize()
+
+    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) ----------------
 
@@ -238,13 +298,19 @@ def _ple_context(self, ids: torch.Tensor, position: int) -> list[int]:
 
     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."""
+        if torch.version.hip is not None:
+            self._await_pending()
+            if self._failed or self._closed:
+                raise RuntimeError("PLE disk table is closed or 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))
+                if self._worker is not None:
+                    return self._start_fill(reqs, bs)
 
                 def _complete() -> None:
                     try:
@@ -281,6 +347,20 @@ def _complete() -> None:
     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)
+        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 and torch.version.hip is not None:
+        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()