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 e33b18a8e0..80860c80d8 100644 --- a/python/freetoken/kernel/csrc/row_store/row_store_ext.cpp +++ b/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 #include @@ -41,6 +41,10 @@ #include +#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM) +#include +#endif + namespace py = pybind11; namespace { @@ -63,6 +67,70 @@ uint8_t *page_aligned_alloc(size_t bytes) { return static_cast(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(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(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); +} + +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(flag.data_ptr()); + release_store_i64(flag.data_ptr(), 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(memop_write(reinterpret_cast(stream), addr, 7)); + if (rc == hipSuccess) rc = static_cast(memop_wait_geq(reinterpret_cast(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(), 0); + rc = hipGraphLaunch(exec, stream); + } + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + const bool ok = rc == hipSuccess && *flag.data_ptr() == 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(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 } diff --git a/python/freetoken/kernel/row_store.py b/python/freetoken/kernel/row_store.py index 5b42479ebc..275c34effa 100644 --- a/python/freetoken/kernel/row_store.py +++ b/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"] diff --git a/python/freetoken/models/deepseek_v41/engram_table.py b/python/freetoken/models/deepseek_v41/engram_table.py index 55450bbd80..6f565e605b 100644 --- a/python/freetoken/models/deepseek_v41/engram_table.py +++ b/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") diff --git a/python/freetoken/models/qwen4_exp/ple_disk.py b/python/freetoken/models/qwen4_exp/ple_disk.py index 536853be3b..b879291343 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,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,6 +298,10 @@ 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) @@ -245,6 +309,8 @@ 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)) + 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: diff --git a/setup.py b/setup.py index 070b8fd5cf..73d30eaa87 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 dec852ec03..6692b89e58 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 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() diff --git a/tests/models/qwen4_exp/test_ple_disk.py b/tests/models/qwen4_exp/test_ple_disk.py index 8381720740..f71c9e6a0f 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