diff --git a/docs/models.md b/docs/models.md index cdadc2dcb1..adfd838e97 100644 --- a/docs/models.md +++ b/docs/models.md @@ -76,6 +76,26 @@ 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 compatibility + +CPU and hybrid MoE have the same execution semantics on CUDA and ROCm, but use +backend-specific GPU/CPU synchronization. CUDA retains its mapped-pinned, +per-slot flag handshake. ROCm 7.14 uses HIP signal memory and explicit Graph +batch-memory-op nodes whose storage follows the CPU executor's lifetime. + +At startup, FreeToken verifies the ROCm path with a real graph capture, +instantiate, and replay probe. If that probe fails, FreeToken disables CUDA +Graph for CPU/hybrid MoE and continues on the correct eager path. Set +`FREETOKEN_CPU_MOE_FLAG_SYNC=0` to explicitly disable the native flag handshake; +on ROCm this also selects the graph-off eager path. + +The same safety check covers `--moe-strategy offload --moe-cpu-layers ...`. +For tensor parallelism, a failed probe on any rank disables Graph on every rank. +On ROCm, `auto` does not upgrade offload to hybrid while Graph is enabled, +because backend selection happens before the executor can run its probe. +ROCm versions without the graph batch-memory-op APIs use the eager fallback; +a version number alone does not enable graph replay. + ## Notes - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 3988caf602..f988fbfc01 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -24,10 +24,18 @@ from freetoken.moe.expert_banks import attach_resident_banks, load_expert_banks from freetoken.moe.host_banks import PinFailed from freetoken.moe.offload_cache import OffloadMoeCache, attach_offload_moe_cache -from freetoken.utils import align_ceil, init_logger, is_sm90_family, is_sm100_family, mem_GB, torch_dtype +from freetoken.utils import ( + align_ceil, + init_logger, + is_rocm, + is_sm90_family, + is_sm100_family, + mem_GB, + torch_dtype, +) from .config import EngineConfig -from .graph import GraphRunner, get_free_memory +from .graph import GraphRunner, _determine_cuda_graph_bs, get_free_memory from .sample import BatchSamplingArgs, Sampler from freetoken.kvcache import create_kv_pool, resolve_pool_class from freetoken.kvcache.base import CacheRebuildRejected @@ -38,6 +46,62 @@ logger = init_logger(__name__) + +def _cuda_graph_disabled(config: EngineConfig) -> bool: + return config.cuda_graph_bs == [] or ( + config.cuda_graph_bs is None and config.cuda_graph_max_bs == 0 + ) + + +def _auto_hybrid_allowed_before_rocm_probe(config: EngineConfig) -> bool: + """Auto cannot verify the HIP replay handshake before constructing an executor.""" + return not is_rocm() or _cuda_graph_disabled(config) + + +def _disable_unsafe_rocm_cpu_moe_graph( + config: EngineConfig, executor, *, tp_cpu_group=None +) -> bool: + """Disable graph capture when ROCm CPU/Hybrid MoE lacks a replay-safe handshake. + + This runs after the executor's real capture/replay capability probe and before + ``GraphRunner`` is constructed. An empty explicit batch list is the canonical graph-off + setting; max_bs is cleared too so later rebuilds cannot silently turn capture back on. + """ + if not is_rocm(): + return False + unsafe = executor is not None and not executor.graph_capture_safe + if tp_cpu_group is not None: + # A local probe failure must disable capture on every rank; mixing eager + # and captured TP collectives can deadlock during startup or replay. + flag = torch.tensor([int(unsafe)], dtype=torch.int32, device="cpu") + torch.distributed.all_reduce( + flag, op=torch.distributed.ReduceOp.MAX, group=tp_cpu_group + ) + unsafe = bool(flag.item()) + if not unsafe or _cuda_graph_disabled(config): + return False + object.__setattr__(config, "cuda_graph_bs", []) + object.__setattr__(config, "cuda_graph_max_bs", 0) + logger.warning_rank0( + "ROCm CPU/Hybrid MoE stream-memory synchronization did not pass capture/replay " + "on at least one rank; " + "disabling CUDA Graph and continuing in the correct eager path" + ) + return True + + +def _cpu_moe_flag_slots_per_layer(config: EngineConfig, free_memory: int) -> int: + """Cover every configured graph batch size while retaining eager headroom.""" + from freetoken.moe.cpu_executor import _FLAG_SLOTS_PER_LAYER + + graph_batch_sizes = _determine_cuda_graph_bs( + cuda_graph_bs=config.cuda_graph_bs, + cuda_graph_max_bs=config.cuda_graph_max_bs, + free_memory=free_memory, + ) + return max(_FLAG_SLOTS_PER_LAYER, len(set(graph_batch_sizes))) + + def _require_offload_cache_size(cache_size: int, num_experts: int) -> None: """The offload MoE cache needs at least one slot per expert per layer. A too-small size (e.g. a bare offload run with moe_cache_size unset and auto disabled) must fail loudly.""" @@ -368,6 +432,7 @@ def __init__(self, config: EngineConfig): self.tp_cpu_group = self._init_communication(config) free_min, free_max = self._sync_get_memory() init_free_memory = free_max # startup KV sizing keeps cross-rank MAX (unchanged) + self._init_free_memory = init_free_memory self._baseline_free = free_min # rebuild baseline: cross-rank MIN, deterministic across ranks logger.info_rank0(f"Free memory before loading model: {mem_GB(init_free_memory)}") @@ -405,6 +470,12 @@ def __init__(self, config: EngineConfig): self._host_tables_bytes = int(self.model.load_host_tables(config) or 0) if is_offload_moe_strategy(config.moe_strategy): self._init_offload_moe_cache(config) + # cpu/hybrid and offload + --moe-cpu-layers all attach the same executor. + # Auto may select hybrid from a bandwidth profile before GPU init; at this point + # it is safe only if the native handshake was verified or graphs are disabled. + _disable_unsafe_rocm_cpu_moe_graph( + config, self.cpu_moe_executor, tp_cpu_group=self.tp_cpu_group + ) if hasattr(self.model, "prepare_for_runtime"): self.model.prepare_for_runtime() self.encoder_cache = None @@ -827,8 +898,20 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: f"(MoE layer {type(sample).__name__} is missing {required})." ) # Decode batches never exceed max_running_req, but CUDA-graph padding can - # round a batch up to the largest captured size; cover both. - max_tokens = max(config.max_running_req, config.cuda_graph_max_bs or 0, 1) + # round a batch up to the largest explicitly captured size. Derive both + # scratch capacity and handshake slots from the same resolved list that + # GraphRunner will use (explicit cuda_graph_bs takes precedence over max_bs). + graph_batch_sizes = _determine_cuda_graph_bs( + cuda_graph_bs=config.cuda_graph_bs, + cuda_graph_max_bs=config.cuda_graph_max_bs, + free_memory=self._init_free_memory, + ) + max_tokens = max( + config.max_running_req, + config.cuda_graph_max_bs or 0, + max(graph_batch_sizes, default=0), + 1, + ) method = sample.quant_method executor = CpuMoeExecutor( cache, @@ -838,6 +921,9 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: num_threads=config.moe_cpu_threads, max_tokens=max_tokens, device=self.device, + flag_slots_per_layer=_cpu_moe_flag_slots_per_layer( + config, self._init_free_memory + ), swiglu_alpha=float(sample.alpha), swiglu_limit=sample.limit, # FIXME: the None branch serves GGUF q4_0 banks, which have no quant method yet; drop it once GGUF joins the quant path @@ -1724,7 +1810,17 @@ def override(attr: str, value: Any): # this is dangerous, use with caution from freetoken.moe.cpu_executor import compiled_extension_supports _act = getattr(model_config, "hidden_act", "silu") - if not _cpu_moe_act_ok: + if not _auto_hybrid_allowed_before_rocm_probe(config): + # The capability probe needs a constructed CPU executor, which auto + # selection does not have yet. Keep auto fail-closed on ROCm when graph + # capture is enabled; explicit cpu/hybrid still probes the native path, + # while graph-off auto configurations may safely use eager hybrid. + logger.info_rank0( + "benchbw profile recommends hybrid, but ROCm CUDA Graph is enabled " + "and the native flag handshake has not been verified yet; staying " + "on offload" + ) + elif not _cpu_moe_act_ok: logger.info_rank0( f"benchbw profile recommends hybrid, but the CPU MoE executor does not " f"support this model's expert activation " 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..886ad76f3c 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -16,11 +16,12 @@ // (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar). #include +#include #include -#include +#include #include +#include #include -#include #include #include #include @@ -676,15 +677,247 @@ float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float #endif // ===================================================================================== -// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or -// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE -// (done[slot]=0 then ready[slot]=1), sync = WAIT_VALUE(done[slot] >= 1). The wait is -// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays -// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, -// 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. +// CUDA/HIP stream memory operations. The GPU side of the flag handshake is: +// submit = WRITE_VALUE(done=0), WRITE_VALUE(ready=slot+1); sync = WAIT_VALUE(done>=1). +// The wait runs on the GPU front end rather than occupying an SM. CUDA uses mapped +// pinned arrays and resolves the driver entry points at runtime. ROCm uses two shared +// 64-bit signals allocated with hipMallocSignalMemory (HIP requires that allocation +// type for WAIT addresses) and writes the task slot into the single ready signal. +// +// HIP's ordinary hipStreamWrite/WaitValue and hipStreamBatchMemOp calls execute while +// a stream is being captured but, as of ROCm 7.14, are not themselves recorded in the +// graph. During capture we therefore add explicit batch-memory-op graph nodes and splice +// them into the stream's current dependency set. The ROCm capability probe below does +// a real capture + instantiate + replay with a host handshake; eager enqueue alone is +// deliberately insufficient to enable the path. + +#if FREETOKEN_USE_ROCM && defined(HIP_VERSION_MAJOR) && defined(HIP_VERSION_MINOR) && \ + ((HIP_VERSION_MAJOR > 7) || (HIP_VERSION_MAJOR == 7 && HIP_VERSION_MINOR >= 14)) +#define FREETOKEN_HAS_HIP_GRAPH_MEMOPS 1 +#else +#define FREETOKEN_HAS_HIP_GRAPH_MEMOPS 0 +#endif + +#if FREETOKEN_USE_ROCM + +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS +static hipError_t hip_add_capture_memop_node( + hipStream_t stream, hipStreamBatchMemOpParams* ops, unsigned int count) { + if (ops == nullptr || count == 0) return hipErrorInvalidValue; + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipGraph_t graph = nullptr; + const hipGraphNode_t* dependencies = nullptr; + size_t dependency_count = 0; + hipError_t rc = hipStreamGetCaptureInfo_v2( + stream, &status, nullptr, &graph, &dependencies, &dependency_count); + if (rc != hipSuccess) return rc; + if (status != hipStreamCaptureStatusActive || graph == nullptr) + return hipErrorIllegalState; + + hipCtx_t context = nullptr; +#if defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#endif + rc = hipCtxGetCurrent(&context); +#if defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + if (rc != hipSuccess) return rc; + + hipBatchMemOpNodeParams params{}; + params.ctx = context; + params.count = count; + params.paramArray = ops; + hipGraphNode_t node = nullptr; + rc = hipGraphAddBatchMemOpNode( + &node, graph, dependencies, dependency_count, ¶ms); + if (rc != hipSuccess) return rc; + return hipStreamUpdateCaptureDependencies( + stream, &node, 1, hipStreamSetCaptureDependencies); +} + +static hipError_t hip_alloc_signal(uint64_t** out) { + void* storage = nullptr; + hipError_t rc = hipExtMallocWithFlags( + &storage, sizeof(uint64_t), hipMallocSignalMemory); + if (rc == hipSuccess) { + *out = static_cast(storage); + __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 hipmemops_probe(uintptr_t /*stream*/, uintptr_t /*scratch_addr*/) { + int device = 0; + int can_wait = 0; + if (hipGetDevice(&device) != hipSuccess || + hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != + hipSuccess || + can_wait == 0) + return false; + + uint64_t* ready = nullptr; + uint64_t* done = nullptr; + hipStream_t probe_stream = nullptr; + hipGraph_t graph = nullptr; + hipGraphExec_t graph_exec = nullptr; + hipError_t rc = hip_alloc_signal(&ready); + if (rc == hipSuccess) rc = hip_alloc_signal(&done); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&probe_stream, hipStreamNonBlocking); + + hipStreamBatchMemOpParams submit_ops[2]{}; + submit_ops[0].writeValue.operation = hipStreamMemOpWriteValue64; + submit_ops[0].writeValue.address = reinterpret_cast(done); + submit_ops[0].writeValue.value64 = 0; + submit_ops[1].writeValue.operation = hipStreamMemOpWriteValue64; + submit_ops[1].writeValue.address = reinterpret_cast(ready); + submit_ops[1].writeValue.value64 = 1; + hipStreamBatchMemOpParams sync_op{}; + sync_op.waitValue.operation = hipStreamMemOpWaitValue64; + sync_op.waitValue.address = reinterpret_cast(done); + sync_op.waitValue.value64 = 1; + sync_op.waitValue.flags = hipStreamWaitValueGte; + + if (rc == hipSuccess) + rc = hipStreamBeginCapture(probe_stream, hipStreamCaptureModeRelaxed); + const bool capture_started = rc == hipSuccess; + if (rc == hipSuccess) + rc = hip_add_capture_memop_node(probe_stream, submit_ops, 2); + if (rc == hipSuccess) + rc = hip_add_capture_memop_node(probe_stream, &sync_op, 1); + // Unsupported beta APIs can fail after BeginCapture; end that capture even + // on failure so its stream/graph resources are not left active during cleanup. + if (capture_started) { + hipError_t end_rc = hipStreamEndCapture(probe_stream, &graph); + if (rc == hipSuccess) rc = end_rc; + } + if (rc == hipSuccess) + rc = hipGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0); + + bool served = false; + std::atomic finished{false}; + std::thread responder; + if (rc == hipSuccess) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + __atomic_store_n(done, 0ULL, __ATOMIC_RELEASE); + responder = std::thread([&] { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (!finished.load(std::memory_order_acquire)) { + if (std::chrono::steady_clock::now() >= deadline) { + // A delayed launch can reset done after the timeout. Keep releasing + // this probe-only wait until the stream has drained. + __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + } else if (__atomic_load_n(ready, __ATOMIC_ACQUIRE) == 1ULL) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + served = true; + break; + } + std::this_thread::yield(); + } + __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + }); + rc = hipGraphLaunch(graph_exec, probe_stream); + if (rc == hipSuccess) rc = hipStreamSynchronize(probe_stream); + } + finished.store(true, std::memory_order_release); + if (responder.joinable()) responder.join(); + const bool ok = rc == hipSuccess && served && + __atomic_load_n(done, __ATOMIC_ACQUIRE) == 1ULL; + + if (graph_exec != nullptr) (void)hipGraphExecDestroy(graph_exec); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (probe_stream != nullptr) (void)hipStreamDestroy(probe_stream); + hip_free_signal(ready); + hip_free_signal(done); + return ok; +} + +static void hipmemop_throw(hipError_t rc, const char* what) { + if (rc != hipSuccess) { + throw std::runtime_error( + std::string("CPU MoE HIP flag handshake: ") + what + " failed: " + + hipGetErrorString(rc)); + } +} + +static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, + int64_t slot, uintptr_t capture_ops_addr = 0) { + auto hip_stream = reinterpret_cast(stream); + auto* done = reinterpret_cast(done_addr); + auto* ready = reinterpret_cast(ready_addr); + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipError_t rc = hipStreamIsCapturing(hip_stream, &status); + hipmemop_throw(rc, "submit capture-status query"); + if (status != hipStreamCaptureStatusActive) { + hipmemop_throw( + hipStreamWriteValue64(hip_stream, done, 0, hipStreamWriteValueDefault), + "hipStreamWriteValue64(done)"); + hipmemop_throw( + hipStreamWriteValue64( + hip_stream, ready, static_cast(slot + 1), + hipStreamWriteValueDefault), + "hipStreamWriteValue64(ready)"); + return; + } + + auto* capture_ops = reinterpret_cast(capture_ops_addr); + if (capture_ops == nullptr) { + throw std::runtime_error( + "CPU MoE HIP graph memops require executor-owned parameter storage"); + } + hipmemop_throw(hip_add_capture_memop_node(hip_stream, capture_ops, 2), + "submit graph batch-memory-op"); +} + +static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t /*slot*/, + uintptr_t capture_ops_addr = 0) { + auto hip_stream = reinterpret_cast(stream); + auto* done = reinterpret_cast(done_addr); + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipError_t rc = hipStreamIsCapturing(hip_stream, &status); + hipmemop_throw(rc, "sync capture-status query"); + if (status != hipStreamCaptureStatusActive) { + hipmemop_throw( + hipStreamWaitValue64( + hip_stream, done, 1, hipStreamWaitValueGte, UINT64_MAX), + "hipStreamWaitValue64(done)"); + return; + } + + auto* capture_ops = reinterpret_cast(capture_ops_addr); + if (capture_ops == nullptr) { + throw std::runtime_error( + "CPU MoE HIP graph memops require executor-owned parameter storage"); + } + hipmemop_throw(hip_add_capture_memop_node(hip_stream, capture_ops, 1), + "sync graph batch-memory-op"); +} + +#else // ROCm headers older than the graph batch-memory-op API + +static bool hipmemops_probe(uintptr_t, uintptr_t) { return false; } +static void cumemop_submit(uintptr_t, uintptr_t, uintptr_t, int64_t, uintptr_t = 0) { + throw std::runtime_error( + "CPU MoE HIP flag handshake requires ROCm 7.14 graph batch-memory-op APIs"); +} +static void cumemop_sync(uintptr_t, uintptr_t, int64_t, uintptr_t = 0) { + throw std::runtime_error( + "CPU MoE HIP flag handshake requires ROCm 7.14 graph batch-memory-op APIs"); +} + +#endif // FREETOKEN_HAS_HIP_GRAPH_MEMOPS + +static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) { + return hipmemops_probe(stream, scratch_addr); +} + +#else // CUDA + #if defined(_WIN32) #include static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } @@ -754,7 +987,7 @@ static void cumemop_check(int rc, const char* what) { } static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, - int64_t slot) { + int64_t slot, uintptr_t /*capture_ops_addr*/ = 0) { auto* s = reinterpret_cast(stream); // Order matters and is preserved by the front end: reset done BEFORE raising ready, // so the coordinator's completion write for THIS step can never be wiped. @@ -766,13 +999,16 @@ static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t read "cuStreamWriteValue64(ready)"); } -static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) { +static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot, + uintptr_t /*capture_ops_addr*/ = 0) { cumemop_check(g_cu_wait64(reinterpret_cast(stream), (unsigned long long)(done_addr + (size_t)slot * 8), 1ULL, kCuWaitValueGeq), "cuStreamWaitValue64(done)"); } +#endif // FREETOKEN_USE_ROCM + struct DotChoice { dot_fn fn; const char* name; @@ -1444,17 +1680,28 @@ 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 one mapped-pinned ready/done entry per task slot. HIP signal memory is + // restricted to one 64-bit word per allocation, so ROCm uses a shared ready command + // (slot+1) and shared done flag. Decode layers are serialized by their stream waits, + // therefore only one CPU task can own that pair at a time. 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 + volatile int64_t* ready_flags = nullptr; + volatile int64_t* done_flags = nullptr; + bool shared_flag_signal = false; +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + uint64_t* owned_ready_signal = nullptr; + uint64_t* owned_done_signal = nullptr; + // ROCm 7.14's beta graph API retains hipBatchMemOpNodeParams::paramArray. + // There are only two immutable node shapes per task slot (submit has two ops; + // sync has one shared op), so allocate one canonical array for each shape once + // and let every rebuilt graph reference it. The vector is sized before capture + // starts and never resized, keeping addresses stable without a module-global, + // append-only allocation. These arrays and the signals now share the executor's + // lifetime, which is also the minimum lifetime required by captured CPU-MoE nodes. + std::vector> hip_graph_submit_ops; + std::array hip_graph_sync_ops{}; +#endif int coord_num_slots = 0; std::vector flag_task; // slot -> task (registered lazily) std::vector flag_served; // slot -> completed dispatch count (tests/debug) @@ -1694,6 +1941,12 @@ struct CpuMoeExecutor { ~CpuMoeExecutor() { coord_stop.store(true); if (coord_thread.joinable()) coord_thread.join(); +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + hip_free_signal(owned_ready_signal); + hip_free_signal(owned_done_signal); + owned_ready_signal = nullptr; + owned_done_signal = nullptr; +#endif { std::lock_guard lk(task_mtx); stop = true; @@ -2175,19 +2428,52 @@ struct CpuMoeExecutor { return (slot >= 0 && slot < static_cast(flag_served.size())) ? flag_served[slot] : 0; } - // 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. - void start_flag_coordinator(uintptr_t ready_ptr, uintptr_t done_ptr, int num_slots, - int pin_core) { - ready_flags = reinterpret_cast(ready_ptr); - done_flags = reinterpret_cast(done_ptr); - coord_num_slots = num_slots; - { - std::lock_guard lk(flag_task_mtx); - if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); + 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_ready_value(int slot) const { + if (ready_flags == nullptr) return 0; + if (shared_flag_signal) { + const int64_t command = flag_load_acquire(&ready_flags[0]); + return command == slot + 1 ? 1 : 0; } - flag_served.assign(num_slots, 0); + return flag_load_acquire(&ready_flags[slot]); + } + + int64_t flag_done_value(int slot) const { + if (done_flags == nullptr) return 0; + return flag_load_acquire(&done_flags[shared_flag_signal ? 0 : slot]); + } + + std::vector pending_flag_slots() const { + std::vector pending; + if (ready_flags == nullptr || done_flags == nullptr) return pending; + if (shared_flag_signal) { + const int64_t command = flag_load_acquire(&ready_flags[0]); + if (command > 0 && command <= coord_num_slots && + flag_load_acquire(&done_flags[0]) == 0) + pending.push_back(static_cast(command - 1)); + return pending; + } + for (int slot = 0; slot < coord_num_slots; ++slot) { + if (flag_load_acquire(&ready_flags[slot]) != 0 && + flag_load_acquire(&done_flags[slot]) == 0) + pending.push_back(slot); + } + return pending; + } + + void poison_flag(int slot) { + if (done_flags != nullptr) + flag_store_release(&done_flags[shared_flag_signal ? 0 : slot], 1); + } + + void launch_flag_coordinator(int pin_core) { coord_stop.store(false); coord_pinned = CPU_MOE_HAS_AFFINITY && pin_core >= 0; coord_thread = std::thread([this, pin_core] { @@ -2203,6 +2489,118 @@ struct CpuMoeExecutor { }); } + // 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. + void start_flag_coordinator(uintptr_t ready_ptr, uintptr_t done_ptr, int num_slots, + int pin_core) { + shared_flag_signal = false; + ready_flags = reinterpret_cast(ready_ptr); + done_flags = reinterpret_cast(done_ptr); + coord_num_slots = num_slots; + { + std::lock_guard lk(flag_task_mtx); + if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); + } + flag_served.assign(num_slots, 0); + launch_flag_coordinator(pin_core); + } + + void start_shared_signal_coordinator(int num_slots, int pin_core) { +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + if (num_slots <= 0) + throw std::invalid_argument("CPU MoE HIP flag slot count must be positive"); + if (owned_ready_signal != nullptr || owned_done_signal != nullptr || + !hip_graph_submit_ops.empty()) + throw std::logic_error( + "CPU MoE HIP flag coordinator is already initialized"); + 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 = nullptr; + owned_done_signal = nullptr; + throw std::runtime_error( + std::string("CPU MoE HIP signal allocation failed: ") + hipGetErrorString(rc)); + } + shared_flag_signal = true; + ready_flags = reinterpret_cast(owned_ready_signal); + done_flags = reinterpret_cast(owned_done_signal); + coord_num_slots = num_slots; + hip_graph_submit_ops.clear(); + hip_graph_submit_ops.resize(num_slots); + for (int slot = 0; slot < num_slots; ++slot) { + auto& ops = hip_graph_submit_ops[slot]; + ops[0].writeValue.operation = hipStreamMemOpWriteValue64; + ops[0].writeValue.address = reinterpret_cast(owned_done_signal); + ops[0].writeValue.value64 = 0; + ops[1].writeValue.operation = hipStreamMemOpWriteValue64; + ops[1].writeValue.address = reinterpret_cast(owned_ready_signal); + ops[1].writeValue.value64 = static_cast(slot + 1); + } + hip_graph_sync_ops[0] = {}; + hip_graph_sync_ops[0].waitValue.operation = hipStreamMemOpWaitValue64; + hip_graph_sync_ops[0].waitValue.address = + reinterpret_cast(owned_done_signal); + hip_graph_sync_ops[0].waitValue.value64 = 1; + hip_graph_sync_ops[0].waitValue.flags = hipStreamWaitValueGte; + { + std::lock_guard lk(flag_task_mtx); + if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); + } + flag_served.assign(num_slots, 0); + launch_flag_coordinator(pin_core); +#else + (void)num_slots; + (void)pin_core; + throw std::runtime_error( + "CPU MoE HIP signal allocation requires ROCm 7.14 graph memops"); +#endif + } + + void submit_flag_memop(uintptr_t stream, uintptr_t done_addr, + uintptr_t ready_addr, int64_t slot) { + uintptr_t capture_ops_addr = 0; +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + if (shared_flag_signal) { + if (slot < 0 || slot >= static_cast(hip_graph_submit_ops.size())) + throw std::out_of_range("CPU MoE HIP flag slot is out of range"); + if (done_addr != reinterpret_cast(owned_done_signal) || + ready_addr != reinterpret_cast(owned_ready_signal)) + throw std::invalid_argument( + "CPU MoE HIP memops must use this executor's signal memory"); + capture_ops_addr = + reinterpret_cast(hip_graph_submit_ops[slot].data()); + } +#endif + cumemop_submit(stream, done_addr, ready_addr, slot, capture_ops_addr); + } + + void sync_flag_memop(uintptr_t stream, uintptr_t done_addr, int64_t slot) { + uintptr_t capture_ops_addr = 0; +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + if (shared_flag_signal) { + if (slot < 0 || slot >= static_cast(hip_graph_submit_ops.size())) + throw std::out_of_range("CPU MoE HIP flag slot is out of range"); + if (done_addr != reinterpret_cast(owned_done_signal)) + throw std::invalid_argument( + "CPU MoE HIP memops must use this executor's signal memory"); + capture_ops_addr = reinterpret_cast(hip_graph_sync_ops.data()); + } +#endif + cumemop_sync(stream, done_addr, slot, capture_ops_addr); + } + + size_t graph_memop_param_count() const { +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + return hip_graph_submit_ops.size() * 2 + + (hip_graph_submit_ops.empty() ? 0 : hip_graph_sync_ops.size()); +#else + return 0; +#endif + } + void coordinator_loop() { // Idle backoff (edge-device friendliness): spin hot only while decode traffic is // flowing, with a TIME-based hot window (an iteration count varies 2-6x with slot @@ -2223,13 +2621,19 @@ 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 begin = shared_flag_signal + ? static_cast(flag_load_acquire(&ready_flags[0]) - 1) + : 0; + const int end = shared_flag_signal ? begin + 1 : coord_num_slots; + for (int L = begin; L < end; ++L) { + if (L < 0 || L >= coord_num_slots) continue; // 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 + volatile int64_t* ready = &ready_flags[shared_flag_signal ? 0 : L]; + if (flag_load_acquire(ready) != 0) { + flag_store_release(ready, 0); // consume this step's doorbell MoeTask* t; { std::lock_guard lk(flag_task_mtx); @@ -2240,7 +2644,7 @@ 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; } @@ -2274,7 +2678,8 @@ 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) { + const int done_count = shared_flag_signal ? 1 : coord_num_slots; + for (int L = 0; L < done_count; ++L) { flag_store_release(&done_flags[L], INT64_MAX); } } @@ -2326,9 +2731,24 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("register_flag_task", &CpuMoeExecutor::register_flag_task, py::arg("slot"), py::arg("task")) .def("flag_served_count", &CpuMoeExecutor::flag_served_count, py::arg("slot")) + .def("flag_ready_address", &CpuMoeExecutor::flag_ready_address) + .def("flag_done_address", &CpuMoeExecutor::flag_done_address) + .def("flag_ready_value", &CpuMoeExecutor::flag_ready_value, py::arg("slot")) + .def("flag_done_value", &CpuMoeExecutor::flag_done_value, py::arg("slot")) + .def("pending_flag_slots", &CpuMoeExecutor::pending_flag_slots) + .def("poison_flag", &CpuMoeExecutor::poison_flag, py::arg("slot")) .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("submit_flag_memop", &CpuMoeExecutor::submit_flag_memop, + py::arg("stream"), py::arg("done_addr"), py::arg("ready_addr"), + py::arg("slot")) + .def("sync_flag_memop", &CpuMoeExecutor::sync_flag_memop, + py::arg("stream"), py::arg("done_addr"), py::arg("slot")) + .def("graph_memop_param_count", &CpuMoeExecutor::graph_memop_param_count) .def("set_input_prequant", [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, py::arg("value")) @@ -2342,10 +2762,23 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { py::arg("enabled")) .def("get_spin_wait", [](const CpuMoeExecutor& e) { return e.spin_wait.load(); }); m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr")); - m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"), - py::arg("ready_addr"), py::arg("slot")); - m.def("memop_sync", &cumemop_sync, py::arg("stream"), py::arg("done_addr"), + // Preserve the established CUDA-facing module API. ROCm graph capture uses the + // executor methods above so its beta node parameter arrays have executor lifetime; + // eager ROCm calls remain valid here as well. + m.def("memop_submit", + [](uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, int64_t slot) { + cumemop_submit(stream, done_addr, ready_addr, slot); + }, + py::arg("stream"), py::arg("done_addr"), py::arg("ready_addr"), py::arg("slot")); + m.def("memop_sync", + [](uintptr_t stream, uintptr_t done_addr, int64_t slot) { + cumemop_sync(stream, done_addr, slot); + }, + py::arg("stream"), py::arg("done_addr"), py::arg("slot")); + m.def("memops_use_shared_signal", []() { + return static_cast(FREETOKEN_HAS_HIP_GRAPH_MEMOPS); + }); // ABI capability marker: the highest ActKind this build implements in the // GENERIC epilogue. CpuMoeExecutor.__init__ probes it before requesting an act // id the epilogue must handle -- a prebuilt .so from before ACT_SWIGLUOAI diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index f519687ed6..d1b0043c56 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -1,18 +1,28 @@ """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) + -> submit flag memop (or eager cudaLaunchHostFunc fallback) + -> sync flag memop (or eager cudaLaunchHostFunc fallback) -> H2D (pinned expert output -> GPU) -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 -the subsequent capture embeds in its host/memcpy nodes. +CUDA and ROCm expose the same ordering and results; only the GPU/CPU handshake +implementation differs. CUDA keeps mapped-pinned ready/done arrays and the +established module-level ``cuStreamWriteValue64``/``cuStreamWaitValue64`` path. +ROCm 7.14 uses executor-owned ``hipMallocSignalMemory`` signals and explicitly +adds batch-memory-op nodes during graph capture. The executor also owns the HIP +node-parameter arrays because captured graphs retain their addresses through +replay. + +Only a ROCm handshake that passes a real capture/instantiate/replay probe is +graph-safe; the Engine disables graph capture and continues eagerly when the +probe fails. 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 the subsequent capture embeds in its host/memcpy nodes. """ from __future__ import annotations @@ -32,26 +42,27 @@ # Flag-based GPU<->CPU handshake for hybrid/cpu decode. The default host-func path # (cudaLaunchHostFunc submit+sync per layer) pays ~30-50us of callback dispatch latency # 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 -# 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 -# watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] + +# on a 75-layer model). Instead the GPU raises a backend-specific "ready" signal at +# submit; a persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets +# the "done" signal the GPU waits on at sync -- no host-func round-trip. CUDA uses +# mapped-pinned per-slot arrays with cuStreamWriteValue64/cuStreamWaitValue64; ROCm uses +# executor-owned HIP signal memory and explicit graph batch-memory-op nodes. 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 a logical slot, so every captured decode graph +# rides the handshake. Where memops are unavailable (Windows WDDM, vGPU, old drivers -- +# functionally probed at startup), CUDA keeps the host-func path (functional, slower); +# ROCm fails closed during capture and the Engine normally selects the eager path. 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 +# Minimum flag slots per MoE layer, retaining eager headroom for small graph sets. The +# Engine raises this to the number of distinct configured graph batch sizes when needed. _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" @@ -169,6 +180,7 @@ def __init__( num_threads: int, max_tokens: int, device: torch.device, + flag_slots_per_layer: int = _FLAG_SLOTS_PER_LAYER, swiglu_alpha: float = 1.702, swiglu_limit: float | None = None, fmt: str | None = None, @@ -214,6 +226,9 @@ def __init__( self.quant_format = fmt self.device = device self.max_tokens = int(max_tokens) + self._flag_slots_per_layer = int(flag_slots_per_layer) + if self._flag_slots_per_layer < 1: + raise ValueError("flag_slots_per_layer must be positive") self.apply_router_weight_on_input = bool(apply_router_weight_on_input) # The per-layer tensors and their pointer tables must outlive the executor # (C++ holds raw addresses into both). @@ -223,7 +238,9 @@ def __init__( ) # Decide the flag handshake up front (env + device + a functional stream-memop - # probe): its coordinator needs a core of its own, which the auto thread sizing + # probe). On ROCm the extension's probe performs a real graph capture + replay, + # not merely an eager enqueue: its coordinator needs a core of its own, which + # the auto thread sizing # below reserves (a coordinator time-slicing against the GEMV workers measurably # destabilizes throughput on fully-subscribed boxes). self._flag_sync = _FLAG_SYNC and device.type == "cuda" @@ -234,11 +251,18 @@ 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.info_rank0( + "cpu-moe flag handshake unavailable: stream memory operations " + "did not pass the capture/replay probe; using " + "cudaLaunchHostFunc only for eager execution" + ) + 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) @@ -296,19 +320,31 @@ def __init__( # self so the coordinator's pinned pointers stay valid for the executor's # lifetime (flag_sync itself was decided above, before thread sizing). self._ready = self._done = self._err = None + self._shared_flag_signal = bool( + self._flag_sync + and getattr(_cpu_moe, "memops_use_shared_signal", lambda: False)() + ) 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 * self._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 self._shared_flag_signal: + # HIP wait addresses must come from hipMallocSignalMemory. The C++ + # executor owns the two exact 64-bit allocations so it can stop/join + # the coordinator before freeing them. + 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 @@ -353,6 +389,11 @@ def __init__( f"spin_wait={self.spin_wait}" ) + @property + def graph_capture_safe(self) -> bool: + """Whether native flag sync is usable (capture/replay-probed on ROCm).""" + return self._flag_sync + def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: """Build a CPU int64 tensor of per-layer base addresses for one bank. @@ -599,9 +640,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 use the current stream. The native handshake orders CPU compute + between the D2H inputs and H2D output, including during graph replay.""" pending = self.decode_submit(layer_id, hidden_states, topk_weights, topk_ids) return self.decode_sync(pending) @@ -620,8 +660,33 @@ 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.""" + # hipLaunchHostFunc is not replay-safe for this CPU handshake on ROCm. The + # Engine normally disables graphs when the native capture/replay probe failed; + # this is the non-bypassable last line of defence for direct executor use or a + # future configuration path that misses that downgrade. + if ( + _IS_ROCM + and not self._flag_sync + and torch.cuda.is_current_stream_capturing() + ): + raise RuntimeError( + "CPU/Hybrid MoE is not CUDA-graph safe on ROCm when " + "stream-memory synchronization is unavailable" + ) + 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 ( + _IS_ROCM + and torch.cuda.is_current_stream_capturing() + and slot is None + ): + raise RuntimeError( + "CPU/Hybrid MoE exhausted its graph-safe ROCm flag slots during " + "CUDA-graph capture" + ) if self._gpu_prequant: # DSV4: apply the reference FP8 round-trip on the GPU (the same kernel the @@ -636,16 +701,24 @@ 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, - ) + stream = torch.cuda.current_stream().cuda_stream + if self._shared_flag_signal: + self._ext.submit_flag_memop( + stream, + self._ext.flag_done_address(), + self._ext.flag_ready_address(), + slot, + ) + else: + # Keep CUDA on the established module-level driver-API path. Only + # ROCm needs executor-owned graph parameter storage. + self._cpu_moe.memop_submit( + stream, self._done.data_ptr(), self._ready.data_ptr(), slot + ) else: stream = torch.cuda.current_stream().cuda_stream self._ext.submit_with_cuda_stream(stream, task) @@ -659,9 +732,13 @@ def decode_sync(self, pending: tuple) -> torch.Tensor: if slot is not None: # 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, - ) + stream = torch.cuda.current_stream().cuda_stream + if self._shared_flag_signal: + self._ext.sync_flag_memop( + stream, self._ext.flag_done_address(), slot + ) + else: + self._cpu_moe.memop_sync(stream, self._done.data_ptr(), slot) else: stream = torch.cuda.current_stream().cuda_stream self._ext.sync_with_cuda_stream(stream, task) @@ -680,12 +757,11 @@ 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()): + pending = set(self._ext.pending_flag_slots()) + 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 +783,7 @@ 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 + self._ext.poison_flag(i) # after err: unblock into a checked failure suspects.pop(i, None) def raise_if_unhealthy(self) -> None: @@ -732,7 +808,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/engine/test_cpu_moe_graph_safety.py b/tests/engine/test_cpu_moe_graph_safety.py new file mode 100644 index 0000000000..c231bdad95 --- /dev/null +++ b/tests/engine/test_cpu_moe_graph_safety.py @@ -0,0 +1,240 @@ +from types import SimpleNamespace + +import pytest + + +@pytest.mark.parametrize("strategy", ["cpu", "hybrid", "offload"]) +def test_rocm_unsafe_cpu_executor_disables_graph(monkeypatch, strategy): + """The same gate covers cpu/hybrid and offload + moe_cpu_layers.""" + import freetoken.engine.engine as engine + + monkeypatch.setattr(engine, "is_rocm", lambda: True) + config = SimpleNamespace( + moe_strategy=strategy, + moe_cpu_layers="0" if strategy == "offload" else None, + cuda_graph_bs=None, + cuda_graph_max_bs=32, + ) + executor = SimpleNamespace(graph_capture_safe=False) + + assert engine._disable_unsafe_rocm_cpu_moe_graph(config, executor) + assert config.cuda_graph_bs == [] + assert config.cuda_graph_max_bs == 0 + + +def test_rocm_verified_cpu_executor_keeps_graph(monkeypatch): + import freetoken.engine.engine as engine + + monkeypatch.setattr(engine, "is_rocm", lambda: True) + config = SimpleNamespace(cuda_graph_bs=[1, 2, 4], cuda_graph_max_bs=4) + + assert not engine._disable_unsafe_rocm_cpu_moe_graph( + config, SimpleNamespace(graph_capture_safe=True) + ) + assert config.cuda_graph_bs == [1, 2, 4] + assert config.cuda_graph_max_bs == 4 + + +def test_non_rocm_keeps_existing_cuda_host_callback_fallback(monkeypatch): + import freetoken.engine.engine as engine + + monkeypatch.setattr(engine, "is_rocm", lambda: False) + config = SimpleNamespace(cuda_graph_bs=None, cuda_graph_max_bs=16) + + assert not engine._disable_unsafe_rocm_cpu_moe_graph( + config, SimpleNamespace(graph_capture_safe=False) + ) + assert config.cuda_graph_bs is None + assert config.cuda_graph_max_bs == 16 + + +@pytest.mark.parametrize( + ("cuda_graph_bs", "cuda_graph_max_bs", "expected"), + [ + (None, 32, False), + ([1, 2], 2, False), + ([1], 0, False), + ([], 32, True), + (None, 0, True), + ], +) +def test_rocm_auto_hybrid_requires_graph_to_be_disabled( + monkeypatch, cuda_graph_bs, cuda_graph_max_bs, expected +): + import freetoken.engine.engine as engine + + monkeypatch.setattr(engine, "is_rocm", lambda: True) + config = SimpleNamespace( + cuda_graph_bs=cuda_graph_bs, cuda_graph_max_bs=cuda_graph_max_bs + ) + + assert engine._auto_hybrid_allowed_before_rocm_probe(config) is expected + + +def test_cpu_moe_flag_slots_cover_every_graph_batch_size(): + import freetoken.engine.engine as engine + + config = SimpleNamespace(cuda_graph_bs=None, cuda_graph_max_bs=128) + + # [1, 2, 4] plus 8..128 in steps of 8: 19 distinct captured sizes. + assert engine._cpu_moe_flag_slots_per_layer(config, free_memory=0) == 19 + + +def test_cpu_moe_flag_slots_keep_eager_headroom_for_small_graph_sets(): + import freetoken.engine.engine as engine + + config = SimpleNamespace(cuda_graph_bs=[1, 2, 2, 4], cuda_graph_max_bs=128) + + assert engine._cpu_moe_flag_slots_per_layer(config, free_memory=0) == 16 + + +@pytest.mark.parametrize("graph_bs, max_bs, tokens, slots", [ + ([1, 129], 16, 129, 16), (None, 128, 128, 19), +]) +def test_cpu_executor_capacity_and_checkpoint_block(monkeypatch, graph_bs, max_bs, tokens, slots): + import torch + import freetoken.engine.engine as engine + import freetoken.moe.cpu_executor as cpu_executor + + constructed = {} + executor = object() + + def build(cache, **kwargs): + constructed.update(kwargs) + return executor + + monkeypatch.setattr(cpu_executor, "CpuMoeExecutor", build) + attached = [] + cache = SimpleNamespace(set_cpu_executor=attached.append) + config = SimpleNamespace( + cuda_graph_bs=graph_bs, cuda_graph_max_bs=max_bs, + max_running_req=32, moe_cpu_threads=2, + ) + layer = SimpleNamespace( + top_k=2, activation="swiglu_clamp", apply_router_weight_on_input=False, + alpha=1.0, limit=7.0, + quant_method=SimpleNamespace(cpu_format="ds_fp4", cfg=SimpleNamespace(act_block=32)), + ) + instance = SimpleNamespace(_init_free_memory=0, device=torch.device("cuda")) + engine.Engine._init_cpu_moe_executor(instance, config, cache, [layer]) + + assert constructed["max_tokens"] == tokens + assert constructed["flag_slots_per_layer"] == slots + assert constructed["fmt"] == "ds_fp4" + assert constructed["act_block"] == 32 + assert attached == [executor] + assert instance.cpu_moe_executor is executor + + +def test_cuda_flag_memops_keep_module_level_api(monkeypatch): + """The ownership refactor must not reroute CUDA through the ROCm-only methods.""" + import freetoken.moe.cpu_executor as cpu_executor + + calls = [] + + class TensorStub: + shape = (1, 8) + + def copy_(self, *args, **kwargs): + return self + + def to(self, *args, **kwargs): + return self + + executor = object.__new__(cpu_executor.CpuMoeExecutor) + executor._flag_sync = True + executor._shared_flag_signal = False + 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( + submit_flag_memop=lambda *args: pytest.fail("ROCm submit used on CUDA"), + sync_flag_memop=lambda *args: pytest.fail("ROCm sync used on CUDA"), + ) + 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_id, bs: 303 + executor._io = {1: io} + + monkeypatch.setattr(cpu_executor, "_IS_ROCM", False) + monkeypatch.setattr( + cpu_executor.torch.cuda, + "current_stream", + lambda: SimpleNamespace(cuda_stream=404), + ) + monkeypatch.setattr(cpu_executor.torch, "empty_like", lambda value: TensorStub()) + + tensor = TensorStub() + pending = executor.decode_submit(0, tensor, tensor, tensor) + executor.decode_sync(pending) + + assert calls == [ + ("submit", (404, 202, 101, 3)), + ("sync", (404, 202, 3)), + ] + + +def test_direct_rocm_capture_without_flag_sync_fails_closed(monkeypatch): + import freetoken.moe.cpu_executor as cpu_executor + + monkeypatch.setattr(cpu_executor, "_IS_ROCM", True) + monkeypatch.setattr( + cpu_executor.torch.cuda, "is_current_stream_capturing", lambda: True + ) + executor = object.__new__(cpu_executor.CpuMoeExecutor) + executor._flag_sync = False + + with pytest.raises(RuntimeError, match="not CUDA-graph safe on ROCm"): + executor.decode_submit(0, None, None, None) + + +def test_rocm_capture_with_exhausted_slots_does_not_use_host_callbacks(monkeypatch): + import freetoken.moe.cpu_executor as cpu_executor + + monkeypatch.setattr(cpu_executor, "_IS_ROCM", True) + monkeypatch.setattr(cpu_executor.torch.cuda, "is_current_stream_capturing", lambda: True) + executor = object.__new__(cpu_executor.CpuMoeExecutor) + executor._flag_sync = True + executor._flag_slots = {} + executor._io_for = lambda bs: {} + executor._task_for = lambda layer, bs: 123 + with pytest.raises(RuntimeError, match="exhausted its graph-safe ROCm flag slots"): + executor.decode_submit(0, SimpleNamespace(shape=(17, 64)), None, None) + + +def _check_rank_consensus(rank, init_method): + from datetime import timedelta + from unittest.mock import patch + + import torch.distributed as dist + import freetoken.engine.engine as engine + + dist.init_process_group( + "gloo", rank=rank, world_size=2, init_method=init_method, + timeout=timedelta(seconds=30), + ) + try: + with patch.object(engine, "is_rocm", return_value=True): + for safe in ((True, True), (True, False), (False, True), (False, False)): + config = SimpleNamespace(cuda_graph_bs=[1, 2], cuda_graph_max_bs=2) + disabled = engine._disable_unsafe_rocm_cpu_moe_graph( + config, SimpleNamespace(graph_capture_safe=safe[rank]), + tp_cpu_group=dist.group.WORLD, + ) + assert disabled == (not all(safe)) + assert config.cuda_graph_bs == ([1, 2] if all(safe) else []) + assert config.cuda_graph_max_bs == (2 if all(safe) else 0) + finally: + dist.destroy_process_group() + + +def test_rocm_probe_failure_disables_graph_on_every_rank(tmp_path): + """Exercise the startup consensus with two real CPU ranks, not two GPUs.""" + import torch.multiprocessing as mp + + mp.spawn(_check_rank_consensus, args=((tmp_path / "gloo-init").as_uri(),), nprocs=2) diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index 547c94c60a..fbce3fd432 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -4,9 +4,9 @@ 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 -routing, i.e. the dependency flows through the pinned buffers, not baked-in). +Part 2 -- CUDA-graph capture/replay: the native flag 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 @@ -468,6 +468,8 @@ def test_cpu_moe_decode_cuda_graph_replay(): w = torch.rand(bs, top_k, device=dev, dtype=torch.float32) # Eager forward: also materializes the per-bs pinned buffers + per-(layer,bs) task. + if torch.version.hip and not ex.graph_capture_safe: + pytest.skip("native ROCm graph handshake unavailable") out_eager = ex.decode(layer, hidden, w, ids) torch.cuda.synchronize() ref_eager = _reference(cache, layer, hidden, ids.cpu(), w.cpu()) @@ -503,15 +505,225 @@ 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" + assert ex._ext.flag_done_value(slot) == 1, "completion flag at rest" + assert ex._ext.flag_ready_value(slot) == 0, "doorbell must be consumed" 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(getattr(torch.version, "hip", None) is None, reason="ROCm guard") +def test_rocm_cpu_moe_capture_without_flag_sync_fails_closed(monkeypatch): + """ROCm must never capture the known-unsafe hipLaunchHostFunc fallback.""" + import freetoken.moe.cpu_executor as cpu_executor + + monkeypatch.setattr(cpu_executor, "_FLAG_SYNC", False) + cache = _make_cache(1, 4, 64, 32) + dev = torch.device("cuda") + stream = torch.cuda.Stream() + torch.cuda.set_stream(stream) + executor = cpu_executor.CpuMoeExecutor( + cache, + top_k=2, + activation="silu", + apply_router_weight_on_input=False, + num_threads=2, + max_tokens=1, + device=dev, + ) + hidden = torch.randn(1, 64, device=dev, dtype=torch.bfloat16) + ids = torch.randint(0, 4, (1, 2), device=dev, dtype=torch.int32) + weights = torch.rand(1, 2, device=dev, dtype=torch.float32) + + with pytest.raises(RuntimeError, match="not CUDA-graph safe on ROCm"): + with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=stream): + hidden.add_(0) # keep capture valid when the executor rejects its own nodes + executor.decode(0, hidden, weights, ids) + + +@pytest.mark.skipif(getattr(torch.version, "hip", None) is None, reason="ROCm guard") +def test_rocm_graph_memop_param_storage_is_bounded_across_rebuilds(): + """Rebuilding graphs must reuse executor-owned HIP batch-op parameters. + + ROCm 7.14 retains each graph node's ``paramArray`` pointer. The executor keeps + one immutable submit array per flag slot plus one shared wait array, so graph + rebuilds must neither invalidate those addresses nor append new storage. + """ + from freetoken.moe.cpu_executor import CpuMoeExecutor + + torch.manual_seed(37) + L, E, H, I, top_k, bs = 1, 4, 64, 32, 2, 1 + cache = _make_cache(L, E, H, I) + stream = torch.cuda.Stream() + torch.cuda.set_stream(stream) + executor = CpuMoeExecutor( + cache, + top_k=top_k, + activation="silu", + apply_router_weight_on_input=False, + num_threads=2, + max_tokens=bs, + device=torch.device("cuda"), + ) + if not executor._flag_sync: + pytest.skip("native ROCm graph handshake unavailable") + + hidden = torch.randn(bs, H, device="cuda", dtype=torch.bfloat16) + ids = torch.randint(0, E, (bs, top_k), device="cuda", dtype=torch.int32) + weights = torch.rand(bs, top_k, device="cuda", dtype=torch.float32) + executor.decode(0, hidden, weights, ids) # materialize the task and flag slot + torch.cuda.synchronize() + + expected_params = 2 * executor._flag_capacity + 1 + assert executor._ext.graph_memop_param_count() == expected_params + for rebuild in range(16): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = executor.decode(0, hidden, weights, ids) + for replay in range(3): + torch.manual_seed(1000 + 3 * rebuild + replay) + hidden.copy_(torch.randn_like(hidden)) + ids.copy_(torch.randint_like(ids, 0, E)) + weights.copy_(torch.rand_like(weights)) + graph.replay() + torch.cuda.synchronize() + ref = _reference(cache, 0, hidden, ids.cpu(), weights.cpu()) + rel = (out.float().cpu() - ref).abs().max() / (ref.abs().max() + 1e-6) + assert rel < 2e-2, f"rebuild={rebuild} replay={replay}: {rel}" + assert executor._ext.graph_memop_param_count() == expected_params + del graph, out + + +@pytest.mark.skipif(getattr(torch.version, "hip", None) is None, reason="ROCm guard") +def test_rocm_graph_memops_support_more_than_sixteen_batch_sizes(): + """Every configured graph size gets a native slot, including size sets > 16.""" + from freetoken.moe.cpu_executor import CpuMoeExecutor + + torch.manual_seed(41) + batch_sizes = list(range(1, 20)) + L, E, H, I, top_k = 1, 4, 64, 32, 2 + cache = _make_cache(L, E, H, I) + stream = torch.cuda.Stream() + torch.cuda.set_stream(stream) + executor = CpuMoeExecutor( + cache, + top_k=top_k, + activation="silu", + apply_router_weight_on_input=False, + num_threads=2, + max_tokens=max(batch_sizes), + device=torch.device("cuda"), + flag_slots_per_layer=len(batch_sizes), + ) + if not executor._flag_sync: + pytest.skip("native ROCm graph handshake unavailable") + + captured = [] + for bs in batch_sizes: + hidden = torch.randn(bs, H, device="cuda", dtype=torch.bfloat16) + ids = torch.randint(0, E, (bs, top_k), device="cuda", dtype=torch.int32) + weights = torch.rand(bs, top_k, device="cuda", dtype=torch.float32) + executor.decode(0, hidden, weights, ids) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = executor.decode(0, hidden, weights, ids) + captured.append((bs, graph, out, hidden, weights, ids)) + + # Keep every graph live, then revisit the oldest slots after the last capture. + for bs, graph, out, hidden, weights, ids in captured: + hidden.copy_(torch.randn_like(hidden)) + ids.copy_(torch.randint_like(ids, 0, E)) + weights.copy_(torch.rand_like(weights)) + graph.replay() + torch.cuda.synchronize() + slot = executor._flag_slots[(0, bs)] + assert executor._ext.flag_served_count(slot) >= 2 + ref = _reference(cache, 0, hidden, ids.cpu(), weights.cpu()) + rel = (out.float().cpu() - ref).abs().max() / (ref.abs().max() + 1e-6) + assert rel < 2e-2, f"bs={bs}: {rel}" + + assert len(executor._flag_slots) == len(batch_sizes) + assert executor._ext.graph_memop_param_count() == 2 * len(batch_sizes) + 1 + + +@pytest.mark.parametrize("overlap", [False, True]) +@pytest.mark.parametrize("fetch", [0, 1, 4]) +def test_hybrid_cpu_gpu_graph_replay(overlap, fetch, monkeypatch): + from freetoken.distributed import set_tp_info, try_get_tp_info + from freetoken.kernel.pinned import alloc_pinned_tensor + import freetoken.layers.moe as moe + from freetoken.layers.quantization import NoQuantConfig + from freetoken.moe.cpu_executor import CpuMoeExecutor + from freetoken.moe.offload_cache import OffloadMoeCache + + if try_get_tp_info() is None: + set_tp_info(rank=0, size=1) + monkeypatch.setattr(moe, "_HYBRID_OVERLAP", overlap) + torch.manual_seed(378) + 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=fetch, + ) + 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) + cpu = CpuMoeExecutor( + cache, top_k=K, activation="silu", apply_router_weight_on_input=False, + num_threads=2, max_tokens=B, device=torch.device("cuda"), + ) + if torch.version.hip and not cpu.graph_capture_safe: + pytest.skip("native ROCm graph handshake unavailable") + cache.set_cpu_executor(cpu) + 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 l, out in enumerate(outputs): + ref = torch.zeros(B, H) + for t in range(B): + for k in range(K): + e = raw_ids[t, k] + gate, up = (banks["gate_up"][l][e].float() @ x[t].float()).chunk(2) + act = (torch.nn.functional.silu(gate) * up).bfloat16().float() + ref[t] += w[t, k] * (banks["down"][l][e].float() @ act) + error = (out.float().cpu() - ref).abs().max() + # GPU gate/up rounds to bf16; CPU GEMV keeps fp32 intermediates. + assert error < 0.03 * ref.abs().max() + 1e-5, (step, l, error) + cpu.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 from freetoken.moe.fused_mxfp4 import run_mxfp4_splitk_decode_experts as _run_mxfp4_splitk_decode_experts @@ -530,6 +742,8 @@ def test_cpu_moe_decode_cuda_graph_replay_mxfp4(): apply_router_weight_on_input=False, num_threads=8, max_tokens=bs, device=dev, swiglu_alpha=alpha, swiglu_limit=limit, ) + if torch.version.hip and not ex.graph_capture_safe: + pytest.skip("native ROCm graph handshake unavailable") b = cache.bank_sources layer_banks = [b[n][layer].to(dev) for n in ( "gate_up_blocks", "gate_up_scales", "gate_up_bias", @@ -588,6 +802,8 @@ def test_cpu_moe_decode_cuda_graph_replay_dsfp4(): apply_router_weight_on_input=False, num_threads=8, max_tokens=bs, device=dev, swiglu_limit=limit, act_block=128, ) + if torch.version.hip and not ex.graph_capture_safe: + pytest.skip("native ROCm graph handshake unavailable") b = cache.bank_sources layer_banks = [b[n][layer].to(dev) for n in ( "gate_up_packed", "gate_up_scale", "down_packed", "down_scale")] diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py index 06f954d8d9..89a972899f 100644 --- a/tests/moe/test_cpu_moe_fp8_block.py +++ b/tests/moe/test_cpu_moe_fp8_block.py @@ -266,3 +266,45 @@ def test_fp16_bits_carry_every_e4m3_byte(): finite = torch.isfinite(ref) assert torch.equal(got[finite], ref[finite]) assert torch.equal(got[~finite].abs(), torch.full((2,), 480.0)) # 0x7F / 0xFF, as e4m3_decode + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA or ROCm") +def test_cpu_fp8_block_graph_replay_tracks_new_routing(): + from freetoken.moe.cpu_executor import CpuMoeExecutor + from freetoken.moe.fused_fp8_block import fused_experts_decode_fp8_block + + L, E, H, I, top_k, bs = 1, 4, 256, 128, 2, 3 + cache = _make_fp8_block_cache(L, E, H, I, seed=37) + executor = CpuMoeExecutor( + cache, top_k=top_k, activation="silu", apply_router_weight_on_input=False, + num_threads=2, max_tokens=bs, device=torch.device("cuda"), + ) + if torch.version.hip and not executor.graph_capture_safe: + pytest.skip("native ROCm graph handshake unavailable") + banks = [cache.bank_sources[name][0].cuda() for name in ( + "gate_up", "gate_up_scale", "down", "down_scale", + )] + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + hidden = torch.randn(bs, H, device="cuda", dtype=torch.bfloat16) + ids = torch.randint(0, E, (bs, top_k), device="cuda", dtype=torch.int32) + weights = torch.rand(bs, top_k, device="cuda") + executor.decode(0, hidden, weights, ids) + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + output = executor.decode(0, hidden, weights, ids) + for replay in range(8): + torch.manual_seed(900 + replay) + hidden.copy_(torch.randn_like(hidden)) + ids.copy_(torch.randint_like(ids, 0, E)) + weights.copy_(torch.rand_like(weights)) + graph.replay() + expected = fused_experts_decode_fp8_block( + hidden, *banks, weights, ids.clone(), "silu", False, + ).float() + stream.synchronize() + rel = (output.float() - expected).abs().max() / (expected.abs().max() + 1e-6) + assert rel < 2e-2, f"replay={replay}: {rel}" + executor.raise_if_unhealthy() diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py index b5ec2f09b4..369f276dd0 100644 --- a/tests/moe/test_cpu_moe_q4_0.py +++ b/tests/moe/test_cpu_moe_q4_0.py @@ -7,8 +7,8 @@ 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 -must recompute from the freshly written pinned routing on each replay). +Part 2 covers CUDA-graph capture/replay (the native flag handshake must recompute +from the freshly written pinned routing on each replay). """ from __future__ import annotations @@ -148,7 +148,7 @@ 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.fused import fused_experts_decode_impl @@ -166,6 +166,8 @@ def test_cpu_moe_decode_q4_0_cuda_graph_replay(): cache, top_k=top_k, activation="gelu_tanh", apply_router_weight_on_input=False, num_threads=8, max_tokens=bs, device=dev, ) + if torch.version.hip and not ex.graph_capture_safe: + pytest.skip("native ROCm graph handshake unavailable") b = cache.bank_sources gate_up_layer = _dequant_bank(b["gate_up"][layer], H, dev) down_layer = _dequant_bank(b["down"][layer], I, dev)