diff --git a/docs/models.md b/docs/models.md index cdadc2dcb1..a9a7e3b3e2 100644 --- a/docs/models.md +++ b/docs/models.md @@ -76,6 +76,20 @@ On GH200 this raises the gather from ~220 to ~410 GB/s (`PCIe-gather` in set so the hybrid split is calibrated against the faster gather. Leave it unset on PCIe GPUs: wider grids add no bandwidth there. +### CPU/Hybrid MoE Graph on ROCm 10 + +CPU and hybrid decode use native `hipStreamWriteValue64` / `hipStreamWaitValue64` +with executor-owned HIP signal memory. ROCm 10 captures these operations into +Graph nodes, so replay waits for the CPU experts to compute the current input +and routing. Startup verifies the handshake with actual capture and replay. +CUDA retains its existing mapped-pinned, per-slot handshake. + +If the probe fails or `FREETOKEN_CPU_MOE_FLAG_SYNC=0` disables it, eager execution +remains available, but Graph capture raises a clear error. Use +`--cuda-graph-max-bs 0` (and remove explicit `--cuda-graph-bs` sizes) for that +configuration. The same guard applies to hybrid and offload with CPU layers. +This path targets ROCm 10; older ROCm versions are outside its support scope. + ## Notes - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 6f85971e9c..82135837cb 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -2,9 +2,8 @@ // // Decode ships activations to the CPU, computes the routed experts here (reading // the pinned host expert banks at full RAM bandwidth), and ships the results -// back. To keep the whole decode path inside a single CUDA graph we expose -// submit/sync as host nodes via cudaLaunchHostFunc -- the callbacks only touch a -// CPU worker pool + pinned host buffers and never call any CUDA API. +// back. ROCm 10 uses native HIP stream memops for graph-safe CPU/GPU ordering; +// host-function callbacks remain the eager fallback on ROCm. // // One task is in flight at a time (per MoE layer): submit() wakes the pool, // sync() blocks the host-func thread until the pool drains. The heavy GEMV runs @@ -684,7 +683,102 @@ float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float // which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max // frequency (GEMV workers -1.5x: the reported edge regression). Availability is // probed functionally at startup (memops_probe); anything unsupported (Windows WDDM, -// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. +// vGPU, old drivers) falls back to cudaLaunchHostFunc, eager-only on ROCm. +#if FREETOKEN_USE_ROCM + +static hipError_t hip_alloc_signal(uint64_t** out) { + void* allocation = nullptr; + hipError_t rc = hipExtMallocWithFlags(&allocation, sizeof(uint64_t), hipMallocSignalMemory); + if (rc == hipSuccess) { + *out = static_cast(allocation); + __atomic_store_n(*out, 0ULL, __ATOMIC_RELEASE); + } + return rc; +} + +static void hip_free_signal(uint64_t* signal) { + if (signal != nullptr) (void)hipFree(signal); +} + +static bool cumemops_probe(uintptr_t, uintptr_t) { + int device = 0, can_wait = 0; + if (hipGetDevice(&device) != hipSuccess || + hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != + hipSuccess || !can_wait) + return false; + + uint64_t* ready = nullptr; + uint64_t* done = nullptr; + hipStream_t stream = nullptr; + hipGraph_t graph = nullptr; + hipGraphExec_t executable = nullptr; + hipError_t rc = hip_alloc_signal(&ready); + if (rc == hipSuccess) rc = hip_alloc_signal(&done); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&stream, hipStreamNonBlocking); + // Keep capture nonblocking even if a runtime enqueues rather than records memops. + if (rc == hipSuccess) __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + if (rc == hipSuccess) rc = hipStreamBeginCapture(stream, hipStreamCaptureModeRelaxed); + const bool capturing = rc == hipSuccess; + if (rc == hipSuccess) rc = hipStreamWriteValue64(stream, ready, 1, hipStreamWriteValueDefault); + if (rc == hipSuccess) rc = hipStreamWaitValue64(stream, done, 1, hipStreamWaitValueGte, UINT64_MAX); + if (capturing) { + hipError_t 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); + if (rc == hipSuccess && nodes != 2) rc = hipErrorNotSupported; + if (rc == hipSuccess) rc = hipGraphInstantiate(&executable, graph, nullptr, nullptr, 0); + + bool served = false; + if (rc == hipSuccess) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + __atomic_store_n(done, 0ULL, __ATOMIC_RELEASE); + 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 probe graph never resets done, so this also releases a delayed launch. + __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + }); + rc = hipGraphLaunch(executable, stream); + if (rc == hipSuccess) rc = hipStreamSynchronize(stream); + responder.join(); + } + const bool ok = rc == hipSuccess && served; + if (executable != nullptr) (void)hipGraphExecDestroy(executable); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (stream != nullptr) (void)hipStreamDestroy(stream); + hip_free_signal(ready); + hip_free_signal(done); + return ok; +} + +static void hipmemop_check(hipError_t rc) { + if (rc != hipSuccess) + throw std::runtime_error(std::string("CPU MoE HIP stream memop failed: ") + hipGetErrorString(rc)); +} + +static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, + int64_t slot) { + auto s = reinterpret_cast(stream); + hipmemop_check(hipStreamWriteValue64(s, reinterpret_cast(done_addr), 0, + hipStreamWriteValueDefault)); + hipmemop_check(hipStreamWriteValue64(s, reinterpret_cast(ready_addr), slot + 1, + hipStreamWriteValueDefault)); +} + +static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t) { + hipmemop_check(hipStreamWaitValue64(reinterpret_cast(stream), + reinterpret_cast(done_addr), 1, + hipStreamWaitValueGte, UINT64_MAX)); +} + +#else + #if defined(_WIN32) #include static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } @@ -773,6 +867,8 @@ static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) { "cuStreamWaitValue64(done)"); } +#endif // FREETOKEN_USE_ROCM + struct DotChoice { dot_fn fn; const char* name; @@ -1444,17 +1540,19 @@ struct CpuMoeExecutor { std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- - // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls - // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU - // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round - // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a - // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side - // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: - // the same pointers are used by the GPU kernels and by this thread). + // CUDA uses mapped-pinned per-slot flags; HIP uses one signal pair with ready=slot+1. + // Stream memops notify this coordinator and wait for its CPU result without a kernel. std::thread coord_thread; std::atomic coord_stop{false}; volatile int64_t* ready_flags = nullptr; // GPU increments, this thread polls volatile int64_t* done_flags = nullptr; // this thread sets, GPU spin-waits + // HIP signals hold one word each. Submit/sync pairs must remain stream-serialized, + // as the executor also shares its IO buffers and worker pool across layers. + bool shared_flag_signal = false; +#if FREETOKEN_USE_ROCM + uint64_t* owned_ready_signal = nullptr; + uint64_t* owned_done_signal = nullptr; +#endif int coord_num_slots = 0; std::vector flag_task; // slot -> task (registered lazily) std::vector flag_served; // slot -> completed dispatch count (tests/debug) @@ -1702,6 +1800,10 @@ struct CpuMoeExecutor { for (auto& th : workers) if (th.joinable()) th.join(); for (MoeTask* t : owned_tasks) delete t; +#if FREETOKEN_USE_ROCM + hip_free_signal(owned_ready_signal); + hip_free_signal(owned_done_signal); +#endif } uintptr_t create_task(int layer_id, int num_tokens, uintptr_t x_ptr, @@ -2175,6 +2277,39 @@ struct CpuMoeExecutor { return (slot >= 0 && slot < static_cast(flag_served.size())) ? flag_served[slot] : 0; } + uintptr_t flag_ready_address() const { return reinterpret_cast(ready_flags); } + uintptr_t flag_done_address() const { return reinterpret_cast(done_flags); } + int64_t flag_done_value() const { return flag_load_acquire(done_flags); } + + std::vector pending_flag_slots() const { + const int64_t command = flag_load_acquire(ready_flags); + if (command > 0 && command <= coord_num_slots && flag_load_acquire(done_flags) == 0) + return {static_cast(command - 1)}; + return {}; + } + + void poison_flag(int) { flag_store_release(done_flags, 1); } + + void start_shared_signal_coordinator(int num_slots, int pin_core) { +#if FREETOKEN_USE_ROCM + if (num_slots <= 0 || owned_ready_signal != nullptr || owned_done_signal != nullptr) + throw std::invalid_argument("CPU MoE HIP coordinator requires fresh signals and positive slots"); + hipError_t rc = hip_alloc_signal(&owned_ready_signal); + if (rc == hipSuccess) rc = hip_alloc_signal(&owned_done_signal); + if (rc != hipSuccess) { + hip_free_signal(owned_ready_signal); + hip_free_signal(owned_done_signal); + owned_ready_signal = owned_done_signal = nullptr; + hipmemop_check(rc); + } + shared_flag_signal = true; + start_flag_coordinator(reinterpret_cast(owned_ready_signal), + reinterpret_cast(owned_done_signal), num_slots, pin_core); +#else + throw std::runtime_error("HIP signal memory requires a ROCm build"); +#endif + } + // Start the busy-poll coordinator over the mapped-pinned flag arrays. ``pin_core`` >= 0 // pins the coordinator to that logical CPU (the worker auto-sizing reserves it), so its // polling never migrates onto / contends with a GEMV worker's core. @@ -2223,13 +2358,14 @@ struct CpuMoeExecutor { bool dozing = false; while (!coord_stop.load(std::memory_order_relaxed)) { bool any = false; - for (int L = 0; L < coord_num_slots; ++L) { + const int first_slot = shared_flag_signal ? static_cast(flag_load_acquire(ready_flags)) - 1 : 0; + for (int L = first_slot; L >= 0 && L < coord_num_slots; ++L) { // Binary handshake (memop-compatible: the GPU-side WAIT compares against an // immediate baked at graph capture, so the protocol resets per step instead of // counting). Acquire: everything the GPU made visible before setting ready -- // the D2H input copies -- is visible to the worker pool after this read. - if (flag_load_acquire(&ready_flags[L]) != 0) { - flag_store_release(&ready_flags[L], 0); // consume this step's doorbell + if (flag_load_acquire(&ready_flags[shared_flag_signal ? 0 : L]) != 0) { + flag_store_release(&ready_flags[shared_flag_signal ? 0 : L], 0); MoeTask* t; { std::lock_guard lk(flag_task_mtx); @@ -2240,10 +2376,11 @@ struct CpuMoeExecutor { sync(coord_pinned && spin_wait.load(std::memory_order_relaxed)); } // Release: the workers' y stores are visible before the GPU sees done. - flag_store_release(&done_flags[L], 1); + flag_store_release(&done_flags[shared_flag_signal ? 0 : L], 1); if (L < static_cast(flag_served.size())) ++flag_served[L]; any = true; } + if (shared_flag_signal) break; } if (any) { last_active = coord_clock::now(); @@ -2274,7 +2411,7 @@ struct CpuMoeExecutor { // caught mid-shutdown exits its sync kernel now instead of owning the watchdog // stall. Runs before the destructor's join() returns, while the flag arrays are // still alive on the Python side. - for (int L = 0; L < coord_num_slots; ++L) { + for (int L = 0; L < (shared_flag_signal ? 1 : coord_num_slots); ++L) { flag_store_release(&done_flags[L], INT64_MAX); } } @@ -2329,6 +2466,13 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, py::arg("ready_ptr"), py::arg("done_ptr"), py::arg("num_slots"), py::arg("pin_core")) + .def("start_shared_signal_coordinator", &CpuMoeExecutor::start_shared_signal_coordinator, + py::arg("num_slots"), py::arg("pin_core")) + .def("flag_ready_address", &CpuMoeExecutor::flag_ready_address) + .def("flag_done_address", &CpuMoeExecutor::flag_done_address) + .def("flag_done_value", &CpuMoeExecutor::flag_done_value) + .def("pending_flag_slots", &CpuMoeExecutor::pending_flag_slots) + .def("poison_flag", &CpuMoeExecutor::poison_flag) .def("set_input_prequant", [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, py::arg("value")) diff --git a/python/freetoken/kernel/triton/moe_align.py b/python/freetoken/kernel/triton/moe_align.py index 55a937426b..dedadbdacf 100644 --- a/python/freetoken/kernel/triton/moe_align.py +++ b/python/freetoken/kernel/triton/moe_align.py @@ -41,6 +41,9 @@ No triton.autotune anywhere in this module: all launch configs below are fixed, chosen from a one-time H100 sweep (see comments at each launch site), matching the upstream (vLLM/sglang) style of hardcoding/heuristics instead of autotuning. + +ROCm uses the staged path with direct atomic counts. This avoids incorrect +histograms and unfilled expert block IDs observed with Triton 3.8 on gfx1201. """ from __future__ import annotations @@ -51,6 +54,8 @@ import triton import triton.language as tl +from freetoken.kernel.backend import is_rocm + _SMALL_CAP = 1024 # fused single-CTA path for numel <= this (covers all decode shapes) @@ -116,6 +121,7 @@ def _fill_and_count( sentinel, # == numel effective_E: tl.constexpr, HIST: tl.constexpr, + USE_HIST: tl.constexpr, BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) @@ -125,14 +131,17 @@ def _fill_and_count( tl.store(sorted_token_ids_ptr + offs, sentinel, mask=offs < sorted_numel) # (b) zero the scatter fill-counter (read in the scatter kernel after a barrier) tl.store(fill_counter_ptr + offs, 0, mask=offs < effective_E) - # (c) per-program register histogram, then one merged atomic per touched bin - # (vs one scattered atomic per element -- far fewer, conflict-free atomics) e = tl.load(topk_ids_ptr + offs, mask=offs < numel, other=-1) valid = (offs < numel) & (e >= 0) & (e < effective_E) - e_h = tl.where(valid, e, HIST - 1) # invalid -> spare top bin (>= effective_E) - h = tl.histogram(e_h, HIST) - le = tl.arange(0, HIST) - tl.atomic_add(counts_ptr + le, h, mask=(le < effective_E) & (h > 0)) + if USE_HIST: + # Merge atomics per bin on CUDA; direct atomics avoid the empty counts + # observed with the ROCm histogram lowering. + e_h = tl.where(valid, e, HIST - 1) + h = tl.histogram(e_h, HIST) + le = tl.arange(0, HIST) + tl.atomic_add(counts_ptr + le, h, mask=(le < effective_E) & (h > 0)) + else: + tl.atomic_add(counts_ptr + e, 1, mask=valid) @triton.jit @@ -238,7 +247,8 @@ def moe_align_block_size( fill_counter = torch.empty((effective_E,), dtype=torch.int32, device=device) cumsum = torch.empty((effective_E + 1,), dtype=torch.int32, device=device) - if 0 < numel <= _SMALL_CAP: + rocm = is_rocm() + if 0 < numel <= _SMALL_CAP and not rocm: # Fixed via H100 sweep (9-config grid; this num_warps ladder -- w2 at # numel<=64, w4@128, w8@256, w16@1024 -- won every decode shape, within 5% # of a live tuned search; forced-fixed beat live-tuned by 25-38% on the @@ -280,6 +290,7 @@ def moe_align_block_size( numel, # sentinel effective_E, triton.next_power_of_2(effective_E + 1), + USE_HIST=not rocm, BLOCK_SIZE=256, num_warps=8, num_stages=3, diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index f519687ed6..193c7ae134 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -1,14 +1,19 @@ """Python wrapper around the ``_cpu_moe`` C++ executor (--moe-strategy cpu). Owns the persistent CPU worker pool, the per-batch-size pinned IO buffers, and -the per-(layer, batch-size) host-func task descriptors. ``decode`` issues the -whole CUDA-graph-capturable sequence on the current stream: +the per-(layer, batch-size) task descriptors. ``decode`` issues this sequence +on the current stream: D2H (hidden, topk_ids, topk_weights -> pinned) - -> submit host node (cudaLaunchHostFunc: enqueue MoE task to the pool) - -> sync host node (cudaLaunchHostFunc: spin until the pool drains) + -> notify the CPU worker pool + -> wait for its result -> H2D (pinned expert output -> GPU) +CUDA uses mapped-pinned per-slot flags. ROCm 10 uses executor-owned HIP signal +memory and native write/wait calls that HIP captures into Graph nodes. Only a +successful capture/replay probe enables the ROCm handshake; its host-function +fallback is eager-only, and direct unsafe capture raises an error. + Buffers and tasks are allocated lazily per batch size. GraphRunner runs an eager ``model.forward()`` at each batch size immediately before capturing it, so the first (eager) call materializes the stable pinned buffers + task pointers that @@ -34,24 +39,23 @@ # per call with the GPU stream idle -- 2 calls per MoE layer per decode step (~6 ms/step # on a 75-layer model). Instead the GPU raises a mapped-pinned "ready" flag at submit; a # persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets a "done" -# flag the GPU waits on at sync -- no host-func round-trip. Both GPU-side operations are -# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64, resolved from the -# driver at runtime): they execute on the GPU front end with no SM-resident kernel, so +# flag the GPU waits on at sync -- no host-func round-trip. CUDA uses mapped-pinned +# per-slot flags; ROCm 10 uses HIP signal memory and native write/wait capture. Both +# execute on the GPU front end with no SM-resident kernel, so # GPU utilization stays truthful during the CPU compute window. (The first cut used a # spin-wait kernel; that pinned reported utilization at 99% and laptop CPU/GPU dynamic # power schedulers responded by clamping the CPU frequency -- a net decode regression on # power-coupled edge devices.) Each (layer, decode batch size) pair gets its own flag # slot, so every captured decode graph rides the handshake. Where memops are unavailable -# (Windows WDDM, vGPU, old drivers -- functionally probed at startup) or the slot -# capacity is exceeded, decode keeps the host-func path (functional, slower). A Python +# (functionally probed at startup) or the slot capacity is exceeded, eager decode +# keeps the host-func path. ROCm capture rejects that unsafe fallback. A Python # watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] + # raise_if_unhealthy) instead of an indefinite stream stall. # Caveat: the coordinator busy-polls one core while decode traffic flows (idle backoff # otherwise); FREETOKEN_CPU_MOE_FLAG_SYNC=0 opts out entirely. _FLAG_SYNC = os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC", "1") != "0" -# Flag slots per MoE layer: covers this many distinct decode batch sizes (captured graph -# sizes plus any eager padded sizes); more than that is unheard of, and the overflow -# just keeps the host-func path for the extra combos. +_IS_ROCM = getattr(torch.version, "hip", None) is not None +# CUDA's existing minimum; ROCm reserves up to max_tokens distinct sizes per layer. _FLAG_SLOTS_PER_LAYER = 16 # Pool threads spin up to 50 ms before they park. FREETOKEN_CPU_MOE_SPIN=0 opts out. _SPIN_WAIT = os.getenv("FREETOKEN_CPU_MOE_SPIN", "1") != "0" @@ -234,11 +238,17 @@ def __init__( if not _cpu_moe.memops_probe( torch.cuda.current_stream().cuda_stream, probe_scratch.data_ptr() ): - logger.info_rank0( - "cpu-moe flag handshake unavailable: CUDA stream memory operations " - "are not supported here (Windows WDDM / vGPU / old driver); using " - "the cudaLaunchHostFunc sync" - ) + if _IS_ROCM: + logger.warning_rank0( + "CPU MoE HIP capture/replay probe failed; only eager execution " + "is safe. Disable Graph with --cuda-graph-max-bs 0." + ) + else: + logger.info_rank0( + "cpu-moe flag handshake unavailable: CUDA stream memory operations " + "are not supported here (Windows WDDM / vGPU / old driver); using " + "the cudaLaunchHostFunc sync" + ) self._flag_sync = False nthreads, core_ids = resolve_threads_and_affinity(num_threads) @@ -297,18 +307,23 @@ def __init__( # lifetime (flag_sync itself was decided above, before thread sizing). self._ready = self._done = self._err = None self._flag_slots: dict[tuple[int, int], int] = {} # (layer_id, bs) -> slot - self._flag_capacity = self.num_layers * _FLAG_SLOTS_PER_LAYER + self._flag_capacity = self.num_layers * ( + max(_FLAG_SLOTS_PER_LAYER, self.max_tokens) if _IS_ROCM else _FLAG_SLOTS_PER_LAYER + ) if self._flag_sync: - self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) self._err = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._ready.zero_() - self._done.zero_() self._err.zero_() - self._ext.start_flag_coordinator( - self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity, - self._coord_core, - ) + if _IS_ROCM: + self._ext.start_shared_signal_coordinator(self._flag_capacity, self._coord_core) + else: + self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) + self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) + self._ready.zero_() + self._done.zero_() + self._ext.start_flag_coordinator( + self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity, + self._coord_core, + ) self._watchdog_stop = False # The thread target holds a WEAKREF and re-derefs it each tick: a bound # method would strong-reference the executor forever (the loop never ends @@ -599,9 +614,8 @@ def decode( ) -> torch.Tensor: """One MoE layer of decode on the CPU. Returns a GPU [bs, H] tensor. - All ops go on the current CUDA stream so the whole thing is captured into - the active CUDA graph (the two host nodes carry the data dependency on the - pinned buffers, which hold this step's real routing on replay).""" + All ops go on the current GPU stream; ROCm capture requires native HIP + flag synchronization so each replay computes this step's routing.""" pending = self.decode_submit(layer_id, hidden_states, topk_weights, topk_ids) return self.decode_sync(pending) @@ -620,8 +634,20 @@ def decode_submit( computes only the routes assigned to it. Returns an opaque handle to pass to :meth:`decode_sync`. The output tensor is allocated here so it stays live (and distinct from the interleaved GPU work) across the overlap window.""" + capturing = ( + _IS_ROCM and self.device.type == "cuda" and torch.cuda.is_current_stream_capturing() + ) + if capturing and not self._flag_sync: + raise RuntimeError( + "CPU/Hybrid MoE Graph capture on ROCm requires the native HIP handshake; " + "disable Graph with --cuda-graph-max-bs 0" + ) bs = hidden_states.shape[0] io = self._io_for(bs) + task = self._task_for(layer_id, bs) + slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None + if capturing and slot is None: + raise RuntimeError("CPU/Hybrid MoE exhausted its native HIP Graph flag slots") if self._gpu_prequant: # DSV4: apply the reference FP8 round-trip on the GPU (the same kernel the @@ -636,15 +662,14 @@ def decode_submit( io["ids"].copy_(topk_ids.to(torch.int32), non_blocking=True) io["w"].copy_(topk_weights.to(torch.float32), non_blocking=True) - task = self._task_for(layer_id, bs) out = torch.empty_like(hidden_states) - slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None if slot is not None: # Front-end memops: done[slot]=0 then ready[slot]=1 (the coordinator's # doorbell). No kernel launched; no host-func round trip. self._cpu_moe.memop_submit( torch.cuda.current_stream().cuda_stream, - self._done.data_ptr(), self._ready.data_ptr(), slot, + self._ext.flag_done_address() if _IS_ROCM else self._done.data_ptr(), + self._ext.flag_ready_address() if _IS_ROCM else self._ready.data_ptr(), slot, ) else: stream = torch.cuda.current_stream().cuda_stream @@ -660,7 +685,8 @@ def decode_sync(self, pending: tuple) -> torch.Tensor: # Front-end WAIT(done[slot] >= 1): blocks this stream's later nodes without # occupying an SM, so GPU utilization stays truthful during the CPU window. self._cpu_moe.memop_sync( - torch.cuda.current_stream().cuda_stream, self._done.data_ptr(), slot, + torch.cuda.current_stream().cuda_stream, + self._ext.flag_done_address() if _IS_ROCM else self._done.data_ptr(), slot, ) else: stream = torch.cuda.current_stream().cuda_stream @@ -680,12 +706,15 @@ def _watchdog_tick(self, suspects: dict) -> None: windows under heavy external load, but a coordinator that made progress in between is alive by definition. ``suspects`` maps slot -> (first_seen, served_at_first_sight) and persists across ticks.""" - stuck = (self._ready == 1) & (self._done == 0) - if not bool(stuck.any()): + if _IS_ROCM: + pending = set(self._ext.pending_flag_slots()) + else: + stuck = (self._ready == 1) & (self._done == 0) + pending = set(stuck.nonzero().flatten().tolist()) + if not pending: suspects.clear() return now = time.monotonic() - pending = set(stuck.nonzero().flatten().tolist()) for slot in list(suspects): if slot not in pending: del suspects[slot] @@ -707,7 +736,10 @@ def _watchdog_tick(self, suspects: dict) -> None: for i in dead: self._err[i] = 1 for i in dead: - self._done[i] = 1 # after err: unblock the stream into a checked failure + if _IS_ROCM: + self._ext.poison_flag(i) + else: + self._done[i] = 1 # after err: unblock the stream into a checked failure suspects.pop(i, None) def raise_if_unhealthy(self) -> None: @@ -732,7 +764,7 @@ def _watchdog_main(executor_ref) -> None: while True: time.sleep(2.0) executor = executor_ref() - if executor is None or executor._watchdog_stop or executor._ready is None: + if executor is None or executor._watchdog_stop or not executor._flag_sync: return try: executor._watchdog_tick(suspects) diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index 547c94c60a..16e85154fc 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -4,13 +4,14 @@ GPU decode kernels (bf16/nvfp4/mxfp4/ds_fp4) on identical banks and routing across batch sizes (fp32-accumulate, so the only spread is reduction order -> tight tol). -Part 2 -- CUDA-graph capture/replay: the cudaLaunchHostFunc submit/sync host nodes -end to end (eager decode, capture, replay with *new* data -> result tracks the new +Part 2 -- GPU graph capture/replay: the CPU/GPU handshake end to end +(eager decode, capture, replay with *new* data -> result tracks the new routing, i.e. the dependency flows through the pinned buffers, not baked-in). """ from __future__ import annotations +import os from types import SimpleNamespace import pytest @@ -441,6 +442,8 @@ def _reference(cache, layer, hidden, ids, w, act="silu", apply_in=False): def test_cpu_moe_decode_cuda_graph_replay(): + if torch.version.hip and os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC") == "0": + pytest.skip("native HIP handshake explicitly disabled") from freetoken.moe.cpu_executor import CpuMoeExecutor torch.manual_seed(0) @@ -503,17 +506,234 @@ def test_cpu_moe_decode_cuda_graph_replay(): assert (layer, bs) in ex._flag_slots, "flag slot expected for the decode task" slot = ex._flag_slots[(layer, bs)] assert ex._ext.flag_served_count(slot) >= 4, "1 eager + 3 replay dispatches expected" - assert int(ex._done[slot]) == 1 and int(ex._ready[slot]) == 0, "handshake at rest" + if torch.version.hip: + assert ex._ext.flag_done_value() == 1 + assert not ex._ext.pending_flag_slots() + else: + assert int(ex._done[slot]) == 1 and int(ex._ready[slot]) == 0, "handshake at rest" assert int(ex._err.sum()) == 0, "watchdog must not fire in normal operation" ex.raise_if_unhealthy() print("cpu moe cuda graph replay OK") +@pytest.mark.skipif(torch.version.hip is None, reason="ROCm native handshake") +@pytest.mark.parametrize("failure", ["disabled", "probe_failed", "slots_exhausted"]) +def test_rocm_graph_rejects_host_callback_fallback(monkeypatch, failure): + import freetoken.moe.cpu_executor as cpu_executor + from freetoken.kernel import _cpu_moe + + monkeypatch.setattr(cpu_executor, "_FLAG_SYNC", failure != "disabled") + if failure == "probe_failed": + monkeypatch.setattr(_cpu_moe, "memops_probe", lambda *args: False) + cache = _make_cache(1, 4, 64, 32) + executor = cpu_executor.CpuMoeExecutor( + cache, top_k=2, activation="silu", apply_router_weight_on_input=False, + num_threads=2, max_tokens=1, device=torch.device("cuda"), + ) + if failure == "slots_exhausted": + assert executor._flag_sync + executor._flag_capacity = 0 + else: + assert not executor._flag_sync + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + x = torch.randn(1, 64, device="cuda", dtype=torch.bfloat16) + ids = torch.randint(4, (1, 2), device="cuda", dtype=torch.int32) + w = torch.rand(1, 2, device="cuda") + eager = executor.decode(0, x, w, ids) + stream.synchronize() + ref = _reference(cache, 0, x, ids.cpu(), w.cpu()) + assert (eager.float().cpu() - ref).abs().max() < 0.02 * ref.abs().max() + message = "flag slots" if failure == "slots_exhausted" else "requires the native HIP handshake" + with pytest.raises(RuntimeError, match=message): + with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=stream): + x.add_(0) + executor.decode(0, x, w, ids) + + +def test_cuda_flag_memops_keep_existing_module_api(monkeypatch): + import freetoken.moe.cpu_executor as cpu_executor + + class TensorStub: + shape = (1, 8) + + def copy_(self, *args, **kwargs): + return self + + def to(self, *args): + return self + + calls = [] + executor = object.__new__(cpu_executor.CpuMoeExecutor) + executor._flag_sync = True + executor._gpu_prequant = False + executor._flag_slots = {(0, 1): 3} + executor._cpu_moe = SimpleNamespace( + memop_submit=lambda *args: calls.append(("submit", args)), + memop_sync=lambda *args: calls.append(("sync", args)), + ) + executor._ext = SimpleNamespace() + executor._ready = SimpleNamespace(data_ptr=lambda: 101) + executor._done = SimpleNamespace(data_ptr=lambda: 202) + io = {name: TensorStub() for name in ("x", "ids", "w", "y")} + executor._io_for = lambda bs: io + executor._task_for = lambda layer, bs: 303 + executor._io = {1: io} + monkeypatch.setattr(cpu_executor, "_IS_ROCM", False) + monkeypatch.setattr(torch.cuda, "current_stream", lambda: SimpleNamespace(cuda_stream=404)) + monkeypatch.setattr(torch, "empty_like", lambda tensor: TensorStub()) + tensor = TensorStub() + executor.decode_sync(executor.decode_submit(0, tensor, tensor, tensor)) + assert calls == [("submit", (404, 202, 101, 3)), ("sync", (404, 202, 3))] + + +def test_rocm_watchdog_releases_shared_signal(monkeypatch): + import freetoken.moe.cpu_executor as cpu_executor + + monkeypatch.setattr(cpu_executor, "_IS_ROCM", True) + poisoned = [] + executor = object.__new__(cpu_executor.CpuMoeExecutor) + executor._ext = SimpleNamespace( + pending_flag_slots=lambda: [0], flag_served_count=lambda slot: 0, + poison_flag=poisoned.append, + ) + executor._err = torch.zeros(1, dtype=torch.int64) + suspects = {} + monkeypatch.setattr(cpu_executor.time, "monotonic", lambda: 1) + executor._watchdog_tick(suspects) + assert not poisoned + monkeypatch.setattr(cpu_executor.time, "monotonic", lambda: 12) + executor._watchdog_tick(suspects) + assert poisoned == [0] + with pytest.raises(RuntimeError, match="watchdog fired"): + executor.raise_if_unhealthy() + + +@pytest.mark.skipif(torch.version.hip is None, reason="ROCm native handshake") +def test_rocm10_graph_replay_with_nineteen_sizes_and_rebuilds(): + from freetoken.moe.cpu_executor import CpuMoeExecutor, _FLAG_SYNC + + if not _FLAG_SYNC: + pytest.skip("native HIP handshake explicitly disabled") + torch.manual_seed(350) + cache = _make_cache(2, 4, 64, 32) + executor = CpuMoeExecutor( + cache, top_k=2, activation="silu", apply_router_weight_on_input=False, + num_threads=2, max_tokens=19, device=torch.device("cuda"), + ) + assert executor._flag_sync + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + captured = [] + for bs in range(1, 20): + x = torch.randn(bs, 64, device="cuda", dtype=torch.bfloat16) + ids = torch.randint(4, (bs, 2), device="cuda", dtype=torch.int32) + w = torch.rand(bs, 2, device="cuda") + for layer in range(2): + executor.decode(layer, x, w, ids) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + outputs = [executor.decode(layer, x, w, ids) for layer in range(2)] + captured.append((graph, outputs, x, ids, w)) + + for rebuild in range(4): + for index, (graph, outputs, x, ids, w) in enumerate(captured): + if rebuild and index == 0: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + outputs = [executor.decode(layer, x, w, ids) for layer in range(2)] + captured[index] = graph, outputs, x, ids, w + x.copy_(torch.randn_like(x)) + ids.copy_(torch.randint_like(ids, 0, 4)) + w.copy_(torch.rand_like(w)) + graph.replay() + stream.synchronize() + for layer, output in enumerate(outputs): + ref = _reference(cache, layer, x, ids.cpu(), w.cpu()) + rel = (output.float().cpu() - ref).abs().max() / ref.abs().max() + assert rel < 0.02, (rebuild, index, layer, rel) + executor.raise_if_unhealthy() + assert len(executor._flag_slots) == 38 + + +@pytest.mark.skipif(torch.version.hip is None, reason="ROCm native handshake") +@pytest.mark.parametrize("overlap", [False, True]) +def test_rocm10_hybrid_graph_replay(overlap, monkeypatch): + import freetoken.layers.moe as moe + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.kernel.pinned import alloc_pinned_tensor + from freetoken.layers.quantization import NoQuantConfig + from freetoken.moe.cpu_executor import CpuMoeExecutor, _FLAG_SYNC + from freetoken.moe.offload_cache import OffloadMoeCache + + if not _FLAG_SYNC: + pytest.skip("native HIP handshake explicitly disabled") + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + monkeypatch.setattr(moe, "_HYBRID_OVERLAP", overlap) + torch.manual_seed(350) + L, E, H, I, K, B = 2, 4, 128, 64, 2, 2 + cache = OffloadMoeCache( + L, E, E, torch.device("cuda"), decode_target="hybrid", hybrid_max_fetch=1, + ) + banks = {} + for name, shape in (("gate_up", (2 * I, H)), ("down", (H, I))): + bank = alloc_pinned_tensor(L * E, *shape, dtype=torch.bfloat16) + bank.copy_(torch.randn(bank.shape) * 0.05) + banks[name] = list(bank.split(E)) + cache.set_bank_sources(banks) + executor = CpuMoeExecutor( + cache, top_k=K, activation="silu", apply_router_weight_on_input=False, + num_threads=2, max_tokens=B, device=torch.device("cuda"), + ) + assert executor._flag_sync + cache.set_cpu_executor(executor) + layers = [moe.OffloadMoELayer(l, E, K, H, I, quant_config=NoQuantConfig()) for l in range(L)] + for layer in layers: + layer.offload_cache = cache + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + hidden = torch.randn(B, H, dtype=torch.bfloat16, device="cuda") + ids = torch.arange(B * K, dtype=torch.int32, device="cuda").reshape(B, K) + weights = torch.rand(B, K, device="cuda") + + def forward(): + return [layer._decode_routed(hidden, weights, ids.clone()) for layer in layers] + + forward() + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + outputs = forward() + for step in range(16): + x = torch.randn(B, H, dtype=torch.bfloat16) + raw_ids = (torch.arange(B * K).reshape(B, K) + step).remainder(E).int() + w = torch.rand(B, K) + w[0, 0] = 0 + hidden.copy_(x) + ids.copy_(raw_ids) + weights.copy_(w) + graph.replay() + stream.synchronize() + for layer, output in enumerate(outputs): + ref = _reference(cache, layer, x, raw_ids, w) + error = (output.float().cpu() - ref).abs().max() + # GPU gate/up rounds to bf16; CPU GEMV retains fp32 intermediates. + assert error < 0.03 * ref.abs().max() + 1e-5, (step, layer, error) + executor.raise_if_unhealthy() + + def test_cpu_moe_decode_cuda_graph_replay_mxfp4(): - """gpt-oss mxfp4 path under capture/replay: the host nodes must recompute the + """gpt-oss mxfp4 path under capture/replay: the handshake must recompute the clamped-swiglu+bias GEMV from the freshly written pinned routing on each replay.""" from freetoken.moe.cpu_executor import CpuMoeExecutor + if torch.version.hip and os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC") == "0": + pytest.skip("native HIP handshake explicitly disabled") from freetoken.moe.fused_mxfp4 import run_mxfp4_splitk_decode_experts as _run_mxfp4_splitk_decode_experts torch.manual_seed(7) @@ -572,6 +792,8 @@ def test_cpu_moe_decode_cuda_graph_replay_dsfp4(): FP8 round-trip -> gate_up -> intermediate FP8 round-trip -> down) must recompute from the freshly written pinned routing on each replay.""" from freetoken.moe.cpu_executor import CpuMoeExecutor + if torch.version.hip and os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC") == "0": + pytest.skip("native HIP handshake explicitly disabled") from freetoken.moe.fused_ds_fp4 import routed_experts_fp4 torch.manual_seed(11) diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py index b5ec2f09b4..731680a470 100644 --- a/tests/moe/test_cpu_moe_q4_0.py +++ b/tests/moe/test_cpu_moe_q4_0.py @@ -7,7 +7,7 @@ production bf16 GPU decode kernel on byte-identical banks: both are W4A16, so the only spread is weight bf16-rounding + reduction order -> tight relative tolerance. -Part 2 covers CUDA-graph capture/replay (the cudaLaunchHostFunc submit/sync nodes +Part 2 covers GPU graph capture/replay (the CPU/GPU submit/sync handshake must recompute from the freshly written pinned routing on each replay). """ @@ -148,9 +148,12 @@ def test_cpu_decode_q4_0_matches_ggml_mmvq(): def test_cpu_moe_decode_q4_0_cuda_graph_replay(): - """Q4_0 CPU path under capture/replay: the host nodes must recompute the GEMV from + """Q4_0 CPU path under capture/replay: the handshake must recompute the GEMV from the freshly written pinned routing on each replay (dep flows through pinned buffers).""" from freetoken.moe.cpu_executor import CpuMoeExecutor + from freetoken.moe.cpu_executor import _FLAG_SYNC + if torch.version.hip and not _FLAG_SYNC: + pytest.skip("native HIP handshake explicitly disabled") from freetoken.moe.fused import fused_experts_decode_impl torch.manual_seed(9) diff --git a/tests/moe/test_fused_moe.py b/tests/moe/test_fused_moe.py index cafe778449..b7abab1062 100644 --- a/tests/moe/test_fused_moe.py +++ b/tests/moe/test_fused_moe.py @@ -2,6 +2,63 @@ import torch +def _check_moe_alignment(topk_ids, sorted_ids, expert_ids, num_padded, block_size, num_experts): + routes = topk_ids.cpu().reshape(-1).long() + counts = torch.bincount(routes, minlength=num_experts) + padded_counts = ((counts + block_size - 1) // block_size) * block_size + expected_experts = torch.repeat_interleave( + torch.arange(num_experts, dtype=torch.int32), padded_counts // block_size + ) + length = int(num_padded.item()) + assert length == int(padded_counts.sum()) + torch.testing.assert_close(expert_ids[: length // block_size].cpu(), expected_experts) + sorted_live = sorted_ids[:length].cpu().long() + assert bool(((sorted_live >= 0) & (sorted_live <= routes.numel())).all()) + valid = sorted_live < routes.numel() + torch.testing.assert_close(torch.sort(sorted_live[valid]).values, torch.arange(routes.numel())) + owners = torch.repeat_interleave(expected_experts.long(), block_size) + torch.testing.assert_close(routes[sorted_live[valid]], owners[valid]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU is required") +@pytest.mark.parametrize("batch_size", [1, 52, 80, 128, 129, 256]) +@pytest.mark.parametrize("pattern", ["skewed", "spread"]) +def test_moe_alignment_preserves_route_ownership(batch_size, pattern): + from freetoken.kernel.triton.moe_align import moe_align_block_size + + num_experts, top_k, block_size = 128, 8, 16 + if pattern == "skewed": + ids = torch.tensor([114, 45, 99, 46, 98, 74, 102, 65], dtype=torch.int32).repeat(batch_size, 1) + else: + ids = (torch.arange(batch_size * top_k, dtype=torch.int32) % num_experts).view(batch_size, top_k) + ids = ids.cuda() + _check_moe_alignment( + ids, *moe_align_block_size(ids, block_size, num_experts), block_size, num_experts + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="GPU is required") +def test_moe_alignment_graph_replay_tracks_changed_experts(): + from freetoken.kernel.triton.moe_align import moe_align_block_size + + ids = torch.tensor( + [114, 45, 99, 46, 98, 74, 102, 65], dtype=torch.int32, device="cuda" + ).repeat(80, 1) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + moe_align_block_size(ids, 16, 128) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + result = moe_align_block_size(ids, 16, 128) + for step in range(5): + routes = torch.arange(ids.numel(), dtype=torch.int32, device="cuda") + ids.copy_(((routes + step * 17) % 128).view_as(ids)) + graph.replay() + _check_moe_alignment(ids, *result, 16, 128) + stream.synchronize() + + def _activation_and_mul(gate_up: torch.Tensor, activation: str) -> torch.Tensor: gate, up = gate_up.chunk(2, dim=-1) if activation == "silu":