diff --git a/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp b/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp index e33b18a8e..109e9bddd 100644 --- a/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp +++ b/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 #include +#include #include #include #include @@ -41,6 +42,10 @@ #include +#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM) +#include +#endif + namespace py = pybind11; namespace { @@ -63,6 +68,109 @@ uint8_t *page_aligned_alloc(size_t bytes) { return static_cast(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(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(stream), + reinterpret_cast(addr), value, hipStreamWriteValueDefault); +} + +int memop_wait_geq(uintptr_t stream, uintptr_t addr, uint64_t value) { + return hipStreamWaitValue64(reinterpret_cast(stream), + reinterpret_cast(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(memop_write( + reinterpret_cast(stream), reinterpret_cast(ready), 1)); + if (rc == hipSuccess) rc = static_cast(memop_wait_geq( + reinterpret_cast(stream), reinterpret_cast(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(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 } diff --git a/python/freetoken/kernel/row_store.py b/python/freetoken/kernel/row_store.py index 5b42479eb..09c74baa6 100644 --- a/python/freetoken/kernel/row_store.py +++ b/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"] diff --git a/python/freetoken/models/deepseek_v41/engram_table.py b/python/freetoken/models/deepseek_v41/engram_table.py index 55450bbd8..3bec75c3f 100644 --- a/python/freetoken/models/deepseek_v41/engram_table.py +++ b/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") diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 536853be3..7ec067bd8 100644 --- a/python/freetoken/models/qwen4_exp/ple_disk.py +++ b/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,7 +284,11 @@ 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) @@ -245,20 +296,13 @@ def host_fill_batch(self, batch: Batch, use_graph: bool): 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: diff --git a/setup.py b/setup.py index 070b8fd5c..73d30eaa8 100644 --- a/setup.py +++ b/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, diff --git a/tests/models/deepseek_v41/test_engram.py b/tests/models/deepseek_v41/test_engram.py index dec852ec0..51a832974 100644 --- a/tests/models/deepseek_v41/test_engram.py +++ b/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() diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py index 838172074..f71c9e6a0 100644 --- a/tests/models/qwen4_exp/test_ple_disk.py +++ b/tests/models/qwen4_exp/test_ple_disk.py @@ -2,6 +2,11 @@ from __future__ import annotations +import os +from pathlib import Path +import subprocess +import sys +import threading from types import SimpleNamespace import pytest @@ -357,3 +362,85 @@ def test_graph_sync_protocol(tmp_path, monkeypatch): gated, _, _ = _make_table(tmp_path) assert not gated._wait_sync assert gated.host_fill_batch(_decode_batch([3, 4], 5), use_graph=True) is None + + +def _wait_sync_large_graph(tmp_path, kernels): + from freetoken.core import Batch, Context, Req, set_global_ctx + from freetoken.distributed import set_tp_info + from freetoken.engine.graph import GraphRunner + + set_tp_info(0, 1) + set_global_ctx(Context(page_size=1)) + disk, oracle, args = _make_table(tmp_path) + emb = _embedding() + row_ids = torch.zeros((8, args.num_ngram_heads), dtype=torch.int64, device="cuda") + pad = torch.zeros(256, device="cuda") + release = threading.Event() + fill = disk.fill + + def delayed_fill(*args, **kwargs): + assert release.wait(timeout=5), "fill was not released" + return fill(*args, **kwargs) + + def forward(): + out = disk.lookup(row_ids) + if torch.cuda.is_current_stream_capturing(): + for _ in range(kernels): + pad.add_(1) + return out + + disk.fill = delayed_fill + req = Req(input_ids=torch.tensor([3, 4, 9], dtype=torch.int32), table_idx=0, + cached_len=2, output_len=3, uid=0, sampling_params=None, cache_handle=None) + attn = SimpleNamespace(init_capture_graph=lambda **kwargs: None, + prepare_for_capture=lambda batch: None, prepare_for_replay=lambda batch: None) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + # Exercise the capture-time upload from #697 as well as the PLE dispatch hook. + runner = GraphRunner(stream=stream, device=torch.device("cuda"), model=SimpleNamespace(forward=forward), + attn_backend=attn, cuda_graph_bs=[8], cuda_graph_max_bs=None, + free_memory=1024**3, max_seq_len=16, + vocab_size=args.num_ngram_heads * disk.head_dim, dummy_req=req) + batch = Batch(reqs=[req] * 8, phase="decode") + runner.pad_batch(batch) + batch.positions = torch.full((8,), 2, dtype=torch.int32, device="cuda") + batch.out_loc = torch.arange(8, dtype=torch.int32, device="cuda") + for step, token in enumerate((9, 13, 5)): + release.clear() + batch.input_ids = torch.full((8,), token, dtype=torch.int32, device="cuda") + with disk.forward_host_ctx(batch, use_graph=True): + # Keep the WAIT unsatisfied when the first replay starts. + assert int(disk._flag[0]) == 0 + timer = threading.Timer(0.05, release.set) + timer.start() + try: + out = runner.replay(batch) + finally: + timer.join() + stream.synchronize() + ids = emb.row_ids(_meta([[token]] * 8, [[3, 4]] * 8, decode=True)).cuda() + assert _bitwise_equal(out.to(disk.dtype), oracle.lookup(ids)) + assert torch.equal(pad, torch.full_like(pad, (step + 1) * kernels)) + assert int(disk._flag[0]) == 0 + runner.destroy_cuda_graphs() + + +@requires_cuda +def test_wait_sync_large_graph(tmp_path): + from freetoken.kernel.row_store import probe_wait_sync + + if not probe_wait_sync("auto", torch.device("cuda")): + pytest.skip("stream memops unavailable") + root = Path(__file__).resolve().parents[3] + env = dict(os.environ, FREETOKEN_PLE_SYNC="wait") + env["PYTHONPATH"] = os.pathsep.join((str(root / "python"), str(root), env.get("PYTHONPATH", ""))) + script = ( + "from pathlib import Path; " + "from tests.models.qwen4_exp.test_ple_disk import _wait_sync_large_graph; " + f"_wait_sync_large_graph(Path({str(tmp_path)!r}), 16384)" + ) + result = subprocess.run( + [sys.executable, "-c", script], cwd=root, env=env, capture_output=True, text=True, timeout=45, + ) + assert result.returncode == 0, result.stdout + result.stderr