diff --git a/docs/install_amd.md b/docs/install_amd.md index ed934ad65..7a228c4ea 100644 --- a/docs/install_amd.md +++ b/docs/install_amd.md @@ -49,7 +49,36 @@ or NVFP4 Marlin/b12x (e.g. `--quant-backend moe.nvfp4=marlin` or ## CPU and hybrid MoE `--moe-strategy cpu`, `hybrid` and `--moe-cpu-layers` sync the GPU with the CPU -expert pool through HIP stream memory operations. ROCm 10 records them into CUDA -graphs itself; on ROCm 7.14, which does not, FreeToken adds the graph nodes -explicitly. When the startup probe of that path fails, CUDA graph capture stops -with an error; run those strategies with `--cuda-graph-max-bs 0`. +expert pool through HIP stream memory operations. As on CUDA, BF16 inputs use +kernel IO when its startup probe passes: a publish kernel copies the inputs to +mapped host memory and notifies the CPU; a combine kernel reads the CPU output. +The publish kernel establishes the memory ordering before notification and +skips CPU work when no route needs it. + +When kernel IO is unavailable or the input is not BF16, FreeToken uses copies. +With the memop handshake, this fallback records a system-release event after +the input copies and before notifying the CPU. The event is recorded as an +external graph node so that every replay retains the publication point. The +kernel IO path does not record this event. If the memop handshake is unavailable, +eager execution uses host callbacks instead. No custom HIP library is required. + +Native memop capture is detected at runtime. If it is unavailable, FreeToken +adds explicit memop graph nodes. Before model capture, a numerical check +compares eager and Graph outputs against synchronous CPU execution, including +changing routes, skipped tasks and Hybrid output accumulation. If kernel IO +fails this check, FreeToken switches to the event + memcpy path and checks it +again. Startup checks only the model's activation dtype; BF16 can use kernel IO, +while FP16 and FP32 use copies. When decode graphs are disabled, the numerical +check runs eagerly without capturing a graph. If the Graph path is unsafe, +startup logs a warning and disables decode graphs for the engine. TP ranks +agree on this decision, so one unsafe rank disables graphs for all ranks. +A failed eager check, or a capture/runtime API error, stops startup rather +than allowing unverified output. + +Hybrid layers whose host banks are OS-locked can also fetch experts through +pinned staging, using the same `--moe-locked-fetch` and +`--moe-locked-max-fetch` settings as CUDA. The staging gather waits through +stream memops when available, or a GPU flag wait otherwise. + +After updating the native CPU MoE sources, rebuild the extension with +`python setup.py build_ext --inplace` in the ROCm environment before serving. diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 590cf80cc..7c9e5102c 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -43,6 +43,17 @@ logger = init_logger(__name__) +def _disable_unsafe_cpu_moe_graphs(config: EngineConfig, executor, dtype, tp_cpu_group) -> None: + safe = torch.tensor([executor.graph_capture_safe(dtype)], dtype=torch.int32) + if config.tp_info.size > 1: + torch.distributed.all_reduce(safe, op=torch.distributed.ReduceOp.MIN, group=tp_cpu_group) + if not safe.item(): + logger.warning_rank0("CPU MoE IO/handshake is not graph-safe on every TP rank for this dtype; disabling decode graphs") + # Clear explicit sizes too; GraphRunner gives them priority over max_bs, including on rebuild. + object.__setattr__(config, "cuda_graph_bs", []) + object.__setattr__(config, "cuda_graph_max_bs", 0) + + def _clamp_decode_batch_to_pools(config: EngineConfig, pool_sizes: list[int]) -> None: """Clamp the decode batch so every GPU-decoded layer's routed experts fit in its pool.""" num_experts = config.model_config.num_experts @@ -1077,6 +1088,11 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: # 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, *(config.cuda_graph_bs or []), 1) max_tokens *= config.decode_rows_per_req + # Match GraphRunner's explicit-size precedence, including an empty list. + if config.cuda_graph_bs is not None: + graph_enabled = bool(config.cuda_graph_bs) + else: + graph_enabled = config.cuda_graph_max_bs is None or config.cuda_graph_max_bs > 0 by_id = sorted(layers, key=lambda layer: layer.layer_id) layer_formats = [layer.quant_method.cpu_spec() for layer in by_id] executor = CpuMoeExecutor( @@ -1087,12 +1103,16 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: num_threads=config.moe_cpu_threads, max_tokens=max_tokens, device=self.device, + activation_dtype=self.dtype, + graph_enabled=graph_enabled, swiglu_alpha=float(sample.alpha), swiglu_limit=sample.limit, layer_formats=layer_formats, ) cache.set_cpu_executor(executor) self.cpu_moe_executor = executor + if graph_enabled: + _disable_unsafe_cpu_moe_graphs(config, executor, self.dtype, self.tp_cpu_group) def _sync_get_memory(self) -> Tuple[int, int]: """Get the min and max free memory across TP ranks.""" diff --git a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp index 395f26841..ad69f5127 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/gpu_sync_hip.cpp @@ -157,13 +157,14 @@ hipError_t wait_done(hipStream_t stream, uint64_t* done, hipStreamBatchMemOpPara struct GpuSync::Impl { uint64_t* ready = nullptr; // signal memory uint64_t* done = nullptr; + std::map publication_events; // Explicit nodes keep these addresses, so one parameter pair per distinct task, never moved (std::map nodes are stable). std::map> submit_ops; // done = 0, ready = task hipStreamBatchMemOpParams wait_op{}; }; -// The memop-free mode serves only the locked fetch, which ROCm keeps off: every GpuSync here needs memops. -GpuSync::GpuSync(bool) : impl_(std::make_unique()) { +// Locked fetch can use the signal words with a kernel wait when memops are unavailable. +GpuSync::GpuSync(bool require_memops) : impl_(std::make_unique()) { hipError_t rc = alloc_signal(&impl_->ready); if (rc == hipSuccess) rc = alloc_signal(&impl_->done); if (rc != hipSuccess) { @@ -172,10 +173,12 @@ GpuSync::GpuSync(bool) : impl_(std::make_unique()) { hip_check(rc, "signal allocation"); } // Measure now: the measurement captures and allocates, which a graph capture in progress would reject. - (void)records_native_memops(); + if (require_memops) (void)records_native_memops(); } GpuSync::~GpuSync() { + for (const auto& [task, event] : impl_->publication_events) + if (event != nullptr) (void)hipEventDestroy(event); free_signal(impl_->ready); free_signal(impl_->done); } @@ -244,6 +247,13 @@ bool GpuSync::probe(uintptr_t /*stream*/) { void GpuSync::gpu_submit(uintptr_t stream, uint64_t task) { auto s = reinterpret_cast(stream); hipStreamBatchMemOpParams* ops = needs_explicit_nodes(s) ? impl_->submit_ops[task].data() : nullptr; + // Keep each task's event alive across graph capture, replay and recapture. + auto& publication = impl_->publication_events[task]; + if (publication == nullptr) + hip_check(hipEventCreateWithFlags(&publication, hipEventDisableTiming | hipEventReleaseToSystem), + "create input publication event"); + // External recording retains the D2H publication point in every graph replay. + hip_check(hipEventRecordWithFlags(publication, s, hipEventRecordExternal), "publish input copies"); // Reset done before raising ready, so the coordinator's completion for this step can never be wiped. hip_check(write_value(s, impl_->done, 0, ops), "write done"); hip_check(write_value(s, impl_->ready, task, ops ? ops + 1 : nullptr), "write ready"); diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index 1f275749d..fea24b40b 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -5,11 +5,12 @@ the whole CUDA-graph-capturable sequence on the current stream: D2H (hidden, topk_ids, topk_weights -> pinned IO rows [:bs]) + -> publish (HIP external event before the memop doorbell) -> submit (doorbell naming (layer, bs), or a cudaLaunchHostFunc host node) -> sync (wait for the CPU pool to finish) -> H2D (pinned expert output -> GPU) -With the doorbell on CUDA, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h). +With kernel IO, a publish kernel does the D2H and the doorbell, and skips the step when no route goes to the CPU; a combine kernel does the H2D (csrc/cpu_moe/gpu_sync_io.h). """ from __future__ import annotations @@ -171,6 +172,8 @@ def __init__( num_threads: int, max_tokens: int, device: torch.device, + activation_dtype: torch.dtype = torch.bfloat16, + graph_enabled: bool = True, swiglu_alpha: float = 1.702, swiglu_limit: float | None = None, layer_formats: list[tuple[str, dict[str, int]]] | None = None, @@ -342,14 +345,12 @@ def __init__( self._kernel_io = False if self._flag_sync: self._ext.start_flag_coordinator(self._coord_core) - # ROCm keeps the copies until the kernel IO has been checked on AMD GPUs. - if torch.version.hip is None: - reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream) - if reason: - logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell") - else: - self._ext.enable_kernel_io() - self._kernel_io = True + reason = _cpu_moe.kernel_io_probe(torch.cuda.current_stream().cuda_stream) + if reason: + logger.warning_rank0(f"cpu-moe kernel IO unavailable ({reason}); using D2H/H2D copies around the doorbell") + else: + self._ext.enable_kernel_io() + self._kernel_io = True 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 @@ -366,6 +367,21 @@ def __init__( ) self._watchdog.start() + self._graph_safe_dtypes = {activation_dtype} + if torch.version.hip is not None: + self._graph_capture_safe = self._flag_sync + if self._flag_sync: + reason = self._probe_graph_io(activation_dtype, capture_graph=graph_enabled) + if reason and self._kernel_io and activation_dtype == torch.bfloat16: + logger.warning_rank0(f"cpu-moe kernel IO probe failed ({reason}); checking event + memcpy IO") + # Python selects the copy path; the native KernelIo allocation lives until executor teardown. + self._kernel_io = False + reason = self._probe_graph_io(activation_dtype, capture_graph=graph_enabled) + if reason: + if not self._eager_io_safe: + raise RuntimeError(f"CPU MoE eager IO probe failed for {activation_dtype}: {reason}") + logger.warning_rank0(f"cpu-moe graph IO probe failed for {activation_dtype} ({reason}); using eager decode") + self._graph_capture_safe = graph_enabled and not reason fmt_label = "; ".join(f"{_format_label(name, opts)}: {n} layers" for (name, opts), n in layer_counts.items()) logger.info_rank0( @@ -375,9 +391,88 @@ def __init__( f"H={self.H} I={self.I} experts={self.num_experts} layers={self.num_layers} " f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens} " f"spin_wait={self.spin_wait} " - f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}" + f"handshake={'kernel' if self._kernel_io else 'event-memop' if self._flag_sync and torch.version.hip is not None else 'memop' if self._flag_sync else 'host-func'}" ) + def graph_capture_safe(self, dtype: torch.dtype) -> bool: + return self._graph_capture_safe and ( + torch.version.hip is None or dtype in self._graph_safe_dtypes + ) + + def _probe_graph_io(self, dtype: torch.dtype, *, capture_graph: bool = True) -> str: + """Qualify IO against synchronous CPU execution, including hybrid accumulation.""" + self._eager_io_safe = False + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + bs = min(2, self.max_tokens) + hidden = torch.empty(3, bs, self.H, device="cuda", dtype=dtype) + ids = torch.empty(3, bs, self.top_k, device="cuda", dtype=torch.int32) + weights = torch.empty(3, bs, self.top_k, device="cuda", dtype=torch.float32) + add = torch.empty_like(hidden) + base = torch.arange(hidden.numel(), dtype=torch.float32).reshape(hidden.shape) + routes = torch.arange(ids.numel()).reshape(ids.shape) + + def inputs(step): + hidden.copy_(((base + step * 7) % 31 - 15) / 64) + cpu_ids = ((routes + step) % self.num_experts).int() + cpu_ids[step % 3].fill_(-1) + cpu_ids[(step + 1) % 3, :, -1] = -1 + ids.copy_(cpu_ids) + weights.copy_(((routes + step) % 7 + 1).float() / (7 * self.top_k)) + add.copy_(((base + step * 3) % 17 - 8) / 32) + + def forward(): + outputs = [] + for i in range(3): + pending = self.decode_submit(i % self.num_layers, hidden[i], weights[i], ids[i]) + outputs.append(self.decode_sync(pending, add_to=add[i].clone() if i == 1 else None)) + return outputs + + def reference(): + # No asynchronous IO or notification: both paths can otherwise share the same bug. + x = hidden + if self._gpu_prequant: + from freetoken.kernel.triton.dsv4.fp8_linear import act_quant_fp8_roundtrip + + x = torch.stack([act_quant_fp8_roundtrip(row, block=self._act_block) for row in hidden]) + host_x, host_ids, host_w, host_add = x.cpu(), ids.cpu(), weights.cpu(), add.cpu() + expected = [] + for i in range(3): + self._io["x"][:bs].copy_(host_x[i]) + self._io["ids"][:bs].copy_(host_ids[i]) + self._io["w"][:bs].copy_(host_w[i]) + self._ext.run(i % self.num_layers, bs) + # Match the IO rounding boundary before hybrid addition; the BF16 combine kernel + # also adds already-rounded BF16 operands in FP32, then rounds back to BF16. + out = self._io["y"][:bs].to(dtype).clone() + expected.append(out + host_add[i] if i == 1 else out) + return expected + + def matches(actual, expected): + return all(torch.equal(got.cpu(), want) for got, want in zip(actual, expected)) + + inputs(0) + eager = forward() + stream.synchronize() + if not matches(eager, reference()): + return "eager output differs from synchronous CPU reference" + self._eager_io_safe = True + if not capture_graph: + return "" + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = forward() + for step in range(1, 4): + inputs(step) + graph.replay() + stream.synchronize() + self.raise_if_unhealthy() + # Read the reference only after replay, so it cannot prefill stale graph IO. + if not matches(captured, reference()): + return "graph output differs from synchronous CPU reference" + return "" + def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: """Build a CPU int64 tensor of per-layer base addresses for one bank. @@ -559,11 +654,10 @@ def decode_submit( bs = hidden_states.shape[0] if bs > self.max_tokens: raise ValueError(f"CPU MoE decode batch {bs} exceeds the executor's max_tokens {self.max_tokens}") - if not self._graph_capture_safe and torch.cuda.is_current_stream_capturing(): + if not self.graph_capture_safe(hidden_states.dtype) and torch.cuda.is_current_stream_capturing(): raise RuntimeError( - "CPU MoE decode cannot be captured into a CUDA graph here: the stream-memop " - "doorbell is off or failed its startup probe, and the host-function sync is not " - "replay-safe on this runtime; run with --cuda-graph-max-bs 0" + "CPU MoE decode cannot be captured into a CUDA graph here: the IO/handshake " + "path is not replay-safe for this runtime and dtype; run with --cuda-graph-max-bs 0" ) if self._gpu_prequant: @@ -592,7 +686,7 @@ def decode_submit( out = torch.empty_like(hidden_states) if self._flag_sync: # Front-end memops: reset done, then ring the doorbell with (layer, bs). - # No kernel launched; no host-func round trip. + # HIP records the input-publication event before these writes. self._ext.flag_signal(stream, layer_id, bs) else: self._ext.submit_with_cuda_stream(stream, layer_id, bs) diff --git a/python/freetoken/moe/locked_fetch.py b/python/freetoken/moe/locked_fetch.py index b2e4331e0..f91770d82 100644 --- a/python/freetoken/moe/locked_fetch.py +++ b/python/freetoken/moe/locked_fetch.py @@ -19,8 +19,8 @@ def locked_fetch_wanted(config) -> bool: - """Whether locked layers fetch through pinned staging: hybrid, --moe-locked-fetch not 0, CUDA (the copy kernels are untested on ROCm).""" - return config.moe_strategy == "hybrid" and config.moe_locked_fetch != 0 and torch.version.hip is None + """Whether hybrid layers fetch through pinned staging instead of staying CPU-only.""" + return config.moe_strategy == "hybrid" and config.moe_locked_fetch != 0 def check_locked_fetch_config(config) -> None: diff --git a/tests/engine/test_moe_cpu_layers.py b/tests/engine/test_moe_cpu_layers.py index feb19300b..2bfee2770 100644 --- a/tests/engine/test_moe_cpu_layers.py +++ b/tests/engine/test_moe_cpu_layers.py @@ -61,6 +61,27 @@ def test_resolve_backend_dispatch(): assert resolve(_cfg("fused", "8"), L) == frozenset() +@pytest.mark.parametrize("local_safe,remote_safe", [(True, True), (False, True), (True, False)]) +def test_cpu_moe_graph_fallback_agrees_across_tp_ranks(monkeypatch, local_safe, remote_safe): + import torch + from freetoken.engine.engine import _disable_unsafe_cpu_moe_graphs + + config = SimpleNamespace(tp_info=SimpleNamespace(size=2), cuda_graph_bs=[1, 4], cuda_graph_max_bs=4) + executor = SimpleNamespace(graph_capture_safe=lambda dtype: local_safe) + group = object() + calls = [] + + def reduce(tensor, *, op, group): + calls.append((op, group)) + tensor.fill_(bool(tensor.item()) and remote_safe) + + monkeypatch.setattr(torch.distributed, "all_reduce", reduce) + _disable_unsafe_cpu_moe_graphs(config, executor, torch.bfloat16, group) + assert calls == [(torch.distributed.ReduceOp.MIN, group)] + assert config.cuda_graph_bs == ([1, 4] if local_safe and remote_safe else []) + assert config.cuda_graph_max_bs == (4 if local_safe and remote_safe else 0) + + if __name__ == "__main__": import sys diff --git a/tests/moe/hip_probe_faults.cpp b/tests/moe/hip_probe_faults.cpp index 4426bf0d1..fe3dfbae1 100644 --- a/tests/moe/hip_probe_faults.cpp +++ b/tests/moe/hip_probe_faults.cpp @@ -34,6 +34,14 @@ extern "C" hipError_t hipStreamWriteValue64(hipStream_t stream, void* address, u return real("hipStreamWriteValue64")(stream, address, value, flags); } +extern "C" hipError_t hipEventRecordWithFlags(hipEvent_t event, hipStream_t stream, unsigned flags) { + if (mode("publication-event") && flags == hipEventRecordExternal) { + ++injected; + return hipErrorNotSupported; + } + return real("hipEventRecordWithFlags")(event, stream, flags); +} + extern "C" hipError_t hipStreamUpdateCaptureDependencies( hipStream_t stream, hipGraphNode_t* dependencies, size_t count, unsigned flags) { if (mode("dependency")) { diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index edd9ddc9d..ed87a3cc7 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -545,6 +545,237 @@ def test_cpu_moe_hip_probe_error_recovery(hip_probe_fault_library, fault): assert result.returncode == 0, result.stdout + result.stderr +@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm") +@pytest.mark.parametrize("fallback", ["handshake-probe", "graph-data"]) +def test_cpu_moe_hip_unsafe_graph_keeps_eager_usable(monkeypatch, fallback): + from freetoken.kernel import _cpu_moe + from freetoken.moe.cpu_executor import CpuMoeExecutor + + if fallback == "handshake-probe": + monkeypatch.setattr(_cpu_moe, "memops_probe", lambda stream: False) + elif fallback == "graph-data": + original = CpuMoeExecutor.decode_sync + + def corrupt_capture(self, *args, **kwargs): + out = original(self, *args, **kwargs) + return out + 1 if torch.cuda.is_current_stream_capturing() else out + + monkeypatch.setattr(CpuMoeExecutor, "decode_sync", corrupt_capture) + + cache = _make_cache(3, 8, 256, 128, scale=0.03) + ex = CpuMoeExecutor( + cache, top_k=4, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=1, device=torch.device("cuda"), + ) + dtype = torch.bfloat16 + assert not ex.graph_capture_safe(dtype) + if fallback == "graph-data": + assert ex._flag_sync and not ex._kernel_io + hidden = torch.randn(1, 256, device="cuda", dtype=torch.bfloat16).to(dtype) + ids = torch.tensor([[0, 2, -1, 5]], device="cuda", dtype=torch.int32) + weights = torch.tensor([[0.2, 0.3, 0.1, 0.4]], device="cuda") + actual = ex.decode(1, hidden, weights, ids) + expected = _reference(cache, 1, hidden, ids.cpu(), weights.cpu()) + assert actual.dtype == dtype + relative = (actual.float().cpu() - expected).abs().max() / (expected.abs().max() + 1e-6) + assert relative < 2e-2, relative + with pytest.raises(RuntimeError, match="--cuda-graph-max-bs 0"): + with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=torch.cuda.Stream()): + ex.decode(1, hidden, weights, ids) + assert torch.isfinite(ex.decode(1, hidden, weights, ids).cpu()).all() + + +@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm") +def test_cpu_moe_hip_event_fallback_when_kernel_io_unavailable(monkeypatch): + from freetoken.kernel import _cpu_moe + from freetoken.moe.cpu_executor import CpuMoeExecutor + + monkeypatch.setattr(_cpu_moe, "kernel_io_probe", lambda stream: "mapped IO unavailable") + cache = _make_cache(3, 8, 256, 128, scale=0.03) + ex = CpuMoeExecutor( + cache, top_k=4, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=1, device=torch.device("cuda"), + ) + assert ex._flag_sync and not ex._kernel_io + assert ex.graph_capture_safe(torch.bfloat16) + + +@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm") +@pytest.mark.parametrize("graph_bs,max_bs", [(None, 0), ([], 4), (None, 1), ([1], 0), (None, None), (None, -1)]) +def test_cpu_moe_hip_kernel_graph_mismatch_uses_event(monkeypatch, caplog, graph_bs, max_bs): + from freetoken.engine.engine import Engine, logger + from freetoken.engine.graph import _determine_cuda_graph_bs + from freetoken.moe.cpu_executor import CpuMoeExecutor + + torch.manual_seed(0) + graph_enabled = bool(_determine_cuda_graph_bs(graph_bs, max_bs, free_memory=1 << 30)) + monkeypatch.setattr(logger, "propagate", True) + original = CpuMoeExecutor.decode_sync + + def corrupt_kernel_capture(self, *args, **kwargs): + out = original(self, *args, **kwargs) + if self._kernel_io and torch.cuda.is_current_stream_capturing(): + return out + 1 + return out + + monkeypatch.setattr(CpuMoeExecutor, "decode_sync", corrupt_kernel_capture) + cache = _make_cache(3, 8, 256, 128, scale=0.03) + cache.set_cpu_executor = lambda executor: setattr(cache, "cpu_executor", executor) + config = SimpleNamespace( + max_running_req=1, decode_rows_per_req=1, moe_cpu_threads=4, + cuda_graph_bs=graph_bs, cuda_graph_max_bs=max_bs, tp_info=SimpleNamespace(size=1), + ) + layers = [SimpleNamespace( + layer_id=i, top_k=4, activation="silu", apply_router_weight_on_input=False, + alpha=1.702, limit=None, quant_method=SimpleNamespace(cpu_spec=lambda: ("bf16", {})), + ) for i in range(3)] + engine = SimpleNamespace(device=torch.device("cuda"), dtype=torch.bfloat16, tp_cpu_group=None) + Engine._init_cpu_moe_executor(engine, config, cache, layers) + ex = engine.cpu_moe_executor + assert "not graph-safe" not in caplog.text + assert bool(_determine_cuda_graph_bs(config.cuda_graph_bs, config.cuda_graph_max_bs, free_memory=1 << 30)) == graph_enabled + assert ex._flag_sync and ex._kernel_io == (not graph_enabled) + assert ex.graph_capture_safe(torch.bfloat16) == graph_enabled + hidden = torch.randn(1, 256, device="cuda", dtype=torch.bfloat16) + ids = torch.tensor([[0, 2, -1, 5]], device="cuda", dtype=torch.int32) + weights = torch.tensor([[0.2, 0.3, 0.1, 0.4]], device="cuda") + if not graph_enabled: + actual = ex.decode(1, hidden, weights, ids) + expected = _reference(cache, 1, hidden, ids.cpu(), weights.cpu()) + relative = (actual.float().cpu() - expected).abs().max() / (expected.abs().max() + 1e-6) + assert relative < 2e-2, relative + return + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = ex.decode(1, hidden, weights, ids) + for step in range(3): + hidden.add_(0.125) + graph.replay() + torch.cuda.synchronize() + expected = _reference(cache, 1, hidden, ids.cpu(), weights.cpu()) + relative = (actual.float().cpu() - expected).abs().max() / (expected.abs().max() + 1e-6) + assert relative < 2e-2, relative + + +@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm") +@pytest.mark.parametrize("kernel_only", [False, True]) +@pytest.mark.parametrize("graph_enabled", [False, True]) +def test_cpu_moe_hip_probe_rejects_common_eager_graph_error(monkeypatch, kernel_only, graph_enabled): + from freetoken.moe.cpu_executor import CpuMoeExecutor + + original = CpuMoeExecutor.decode_sync + + def corrupt(self, pending, add_to=None): + out = original(self, pending, add_to=add_to) + return out + 1 if not kernel_only or pending[2] else out + + monkeypatch.setattr(CpuMoeExecutor, "decode_sync", corrupt) + cache = _make_cache(3, 8, 256, 128, scale=0.03) + + def make(): + return CpuMoeExecutor(cache, top_k=4, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=2, device=torch.device("cuda"), graph_enabled=graph_enabled) + + if kernel_only: + ex = make() + assert not ex._kernel_io + assert ex.graph_capture_safe(torch.bfloat16) == graph_enabled + else: + with pytest.raises(RuntimeError, match="eager IO probe failed"): + make() + + +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16, torch.float32]) +@pytest.mark.parametrize("accumulate", [False, True]) +def test_cpu_moe_graph_dtype_and_hybrid_accumulation(dtype, accumulate): + from freetoken.moe.cpu_executor import CpuMoeExecutor + + cache = _make_cache(3, 8, 256, 128, scale=0.03) + ex = CpuMoeExecutor(cache, top_k=4, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=2, device=torch.device("cuda"), activation_dtype=dtype) + assert ex.graph_capture_safe(dtype) + hidden = torch.randn(2, 256, device="cuda", dtype=dtype) + ids = torch.tensor([[0, 2, -1, 5], [-1, -1, -1, -1]], device="cuda", dtype=torch.int32) + weights = torch.rand(2, 4, device="cuda") + add = torch.randn_like(hidden) + + def forward(): + pending = ex.decode_submit(1, hidden, weights, ids) + assert pending[2] == (ex._kernel_io and dtype == torch.bfloat16) + return ex.decode_sync(pending, add_to=add.clone() if accumulate else None) + + forward() + torch.cuda.synchronize() + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + actual = forward() + for step in range(6): + hidden.add_(0.125) + ids.fill_(-1) if step % 2 else ids[0].copy_(torch.tensor([1, 3, -1, 7], device="cuda")) + graph.replay() + torch.cuda.synchronize() + ex._io["x"][:2].copy_(hidden.cpu()) + ex._io["ids"][:2].copy_(ids.cpu()) + ex._io["w"][:2].copy_(weights.cpu()) + ex._ext.run(1, 2) + expected = ex._io["y"][:2].to(dtype).clone() + if accumulate: + expected += add.cpu() + assert torch.equal(actual.cpu(), expected) + + +@pytest.mark.parametrize("kernel_io", [False, True]) +def test_cpu_moe_hip_event_record_failure_is_loud(hip_probe_fault_library, kernel_io): + import os + from pathlib import Path + import subprocess + import sys + + script = r""" +import ctypes, os, runpy, sys +import pytest +import torch +from freetoken.kernel import _cpu_moe +from freetoken.moe.cpu_executor import CpuMoeExecutor +helpers = runpy.run_path(sys.argv[1]) +shim = ctypes.CDLL(sys.argv[2]) +kernel_io = sys.argv[3] == 'True' +if not kernel_io: + _cpu_moe.kernel_io_probe = lambda stream: 'forced memcpy fallback' +cache = helpers['_make_cache'](3, 8, 256, 128, scale=0.03) +def make_executor(): + return CpuMoeExecutor(cache, top_k=4, activation='silu', apply_router_weight_on_input=False, + num_threads=4, max_tokens=1, device=torch.device('cuda')) +if kernel_io: + ex = make_executor() + assert ex._kernel_io and ex.graph_capture_safe(torch.bfloat16) + hidden = torch.randn(1, 256, device='cuda', dtype=torch.bfloat16) + ids = torch.tensor([[0, 2, -1, 5]], device='cuda', dtype=torch.int32) + weights = torch.tensor([[0.2, 0.3, 0.1, 0.4]], device='cuda') + ex.decode(1, hidden, weights, ids) + torch.cuda.synchronize() + assert shim.probe_test_injected() == 0 + with pytest.raises(RuntimeError, match='publish input copies'): + ex.decode(1, hidden.half(), weights, ids) +else: + with pytest.raises(RuntimeError, match='publish input copies'): + make_executor() +torch.cuda.synchronize() +assert shim.probe_test_injected() > 0 +""" + env = dict(os.environ, LD_PRELOAD=str(hip_probe_fault_library), + FREETOKEN_TEST_HIP_PROBE_FAULT="publication-event", FREETOKEN_CPU_MOE_FLAG_SYNC="1") + result = subprocess.run( + [sys.executable, "-c", script, str(Path(__file__).resolve()), str(hip_probe_fault_library), str(kernel_io)], + env=env, capture_output=True, text=True, timeout=30, + ) + assert result.returncode == 0, result.stdout + result.stderr + + def test_cpu_moe_decode_cuda_graph_replay(): from freetoken.moe.cpu_executor import CpuMoeExecutor diff --git a/tests/moe/test_cpu_moe_fp8_block.py b/tests/moe/test_cpu_moe_fp8_block.py index baecb86f9..17cc5fb57 100644 --- a/tests/moe/test_cpu_moe_fp8_block.py +++ b/tests/moe/test_cpu_moe_fp8_block.py @@ -261,3 +261,108 @@ 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") +@pytest.mark.parametrize("bs", [1, 3]) +@pytest.mark.parametrize("accumulate", [False, True]) +def test_cpu_fp8_block_graph_distinct_layer_inputs(bs, accumulate): + from freetoken.moe.cpu_executor import CpuMoeExecutor + + L, E, H, I, top_k = 3, 8, 256, 128, 4 + cache = _make_fp8_block_cache(L, E, H, I, seed=7) + stream = torch.cuda.Stream() + with torch.cuda.stream(stream): + ex = CpuMoeExecutor( + cache, top_k=top_k, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=bs, device=torch.device("cuda"), + ) + hidden = torch.randn(L, bs, H, device="cuda", dtype=torch.bfloat16) + ids = torch.zeros(L, bs, top_k, device="cuda", dtype=torch.int32) + weights = torch.rand(L, bs, top_k, device="cuda") + gpu_out = torch.randn_like(hidden) + + def forward(): + outputs = [] + for layer in range(L): + pending = ex.decode_submit(layer, hidden[layer], weights[layer], ids[layer]) + outputs.append(ex.decode_sync(pending, add_to=gpu_out[layer].clone() if accumulate else None)) + return outputs + + forward() + stream.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + captured = forward() + + for step in range(3): + torch.manual_seed(101 + step) + hidden.copy_(torch.randn(L, bs, H, dtype=torch.bfloat16)) + new_ids = torch.randint(0, E, (L, bs, top_k), dtype=torch.int32) + new_ids[step].fill_(-1) + new_ids[(step + 1) % L, :, 0] = -1 + ids.copy_(new_ids) + weights.copy_(torch.rand(L, bs, top_k)) + gpu_out.copy_(torch.randn(L, bs, H, dtype=torch.bfloat16)) + # Do not run eager first: it would prefill the shared host IO with this step's data. + graph.replay() + stream.synchronize() + eager = forward() + stream.synchronize() + for got, expected in zip(captured, eager): + torch.testing.assert_close(got, expected, rtol=0, atol=0) + skipped = gpu_out[step] if accumulate else torch.zeros_like(gpu_out[step]) + torch.testing.assert_close(captured[step], skipped, rtol=0, atol=0) + + +@pytest.mark.skipif(torch.version.hip is None, reason="requires ROCm") +@pytest.mark.parametrize("kernel_io", [True, False]) +def test_cpu_fp8_block_graph_replay_and_recapture(monkeypatch, kernel_io): + from freetoken.kernel import _cpu_moe + from freetoken.moe.cpu_executor import CpuMoeExecutor + + if not kernel_io: + monkeypatch.setattr(_cpu_moe, "kernel_io_probe", lambda stream: "forced memcpy fallback") + L, E, H, I, top_k = 3, 8, 256, 128, 4 + cache = _make_fp8_block_cache(L, E, H, I, seed=71) + ex = CpuMoeExecutor( + cache, top_k=top_k, activation="silu", apply_router_weight_on_input=False, + num_threads=4, max_tokens=3, device=torch.device("cuda"), + ) + assert ex._flag_sync and ex._kernel_io == kernel_io + assert ex.graph_capture_safe(torch.bfloat16) + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + # Reuse the executor after destroying each graph, including a repeated task size. + for bs in (1, 3, 1): + hidden = torch.empty(L, bs, H, device="cuda", dtype=torch.bfloat16) + ids = torch.empty(L, bs, top_k, device="cuda", dtype=torch.int32) + weights = torch.empty(L, bs, top_k, device="cuda", dtype=torch.float32) + + def inputs(step): + generator = torch.Generator().manual_seed(1700 + step) + hidden.copy_(torch.randn(hidden.shape, dtype=hidden.dtype, generator=generator)) + routes = torch.randint(E, ids.shape, dtype=ids.dtype, generator=generator) + routes[step % L].fill_(-1) + ids.copy_(routes) + weights.copy_(torch.rand(weights.shape, generator=generator)) + + def forward(): + return [ex.decode(layer, hidden[layer], weights[layer], ids[layer]) for layer in range(L)] + + inputs(0) + # The memcpy fallback may allocate its publication event inside capture. + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + outputs = forward() + for step in range(1, 11): + inputs(step) + graph.replay() + stream.synchronize() + references = forward() + stream.synchronize() + for actual, expected in zip(outputs, references): + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + graph.reset() + ex.raise_if_unhealthy() diff --git a/tests/moe/test_offload.py b/tests/moe/test_offload.py index 3baddeb52..92ea01ee3 100644 --- a/tests/moe/test_offload.py +++ b/tests/moe/test_offload.py @@ -1129,14 +1129,18 @@ def test_locked_hybrid_layer_fetches_through_staging_in_a_graph(monkeypatch, mem # memops: the GPU waits on the doorbell with a stream memop, else the gather kernel spins on it from types import SimpleNamespace - from freetoken.moe.locked_fetch import LockedLayerFetch + from freetoken.moe.locked_fetch import attach_locked_fetch, settle_banks monkeypatch.setenv("FREETOKEN_LOCKED_FETCH_THREADS", "2") torch.manual_seed(0) cache, sources = _make_split_cache(num_layers=2, locked=(1,), device="cuda") cache.decode_target = "hybrid" cache.locked_hybrid = frozenset({1}) - cache.locked_fetch = LockedLayerFetch(cache, [1], plan={1: (0, 1.0)}, max_fetch=2, executor=SimpleNamespace(core_ids=[], _coord_core=-1, _flag_sync=memops)) + config = SimpleNamespace(moe_strategy="hybrid", moe_locked_fetch=1.0, moe_locked_max_fetch=2) + _, _, staging = settle_banks(config, SimpleNamespace(sources=sources), probe=False, has_locked=True) + assert staging is not None + cache.cpu_executor = SimpleNamespace(core_ids=[], _coord_core=-1, _flag_sync=memops) + attach_locked_fetch(cache, config, staging, layer_rows=None) top_k = 3 ids = torch.zeros(top_k, dtype=torch.int32, device="cuda")