diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index b93de55c2..9b8bdaacb 100644 --- a/python/freetoken/engine/engine.py +++ b/python/freetoken/engine/engine.py @@ -40,6 +40,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 @@ -963,6 +974,7 @@ def _init_cpu_moe_executor(self, config: EngineConfig, cache, layers) -> None: ) cache.set_cpu_executor(executor) self.cpu_moe_executor = executor + _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/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index 1f275749d..8f4585516 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -9,7 +9,7 @@ -> 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 @@ -342,14 +342,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 +364,14 @@ def __init__( ) self._watchdog.start() + if torch.version.hip is not None: + # The flag probe alone misses stale shared IO in HIP graphs using memcpy + memops. + self._graph_capture_safe = self._flag_sync and self._kernel_io + if self._graph_capture_safe: + reason = self._probe_kernel_graph_io() + if reason: + self._graph_capture_safe = False + logger.warning_rank0(f"cpu-moe kernel IO graph probe failed ({reason}); using eager decode") fmt_label = "; ".join(f"{_format_label(name, opts)}: {n} layers" for (name, opts), n in layer_counts.items()) logger.info_rank0( @@ -378,6 +384,51 @@ def __init__( f"handshake={'kernel' if self._kernel_io else 'memop' if self._flag_sync else 'host-func'}" ) + def graph_capture_safe(self, dtype: torch.dtype) -> bool: + # Non-bf16 inputs still take the memcpy path, even when kernel IO is enabled. + return self._graph_capture_safe and ( + torch.version.hip is None or (self._kernel_io and dtype == torch.bfloat16) + ) + + def _probe_kernel_graph_io(self) -> str: + """Check distinct tasks sharing IO, including CPU-route skips, before engine capture.""" + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + hidden = torch.empty(3, 1, self.H, device="cuda", dtype=torch.bfloat16) + ids = torch.empty(3, 1, self.top_k, device="cuda", dtype=torch.int32) + weights = torch.empty(3, 1, self.top_k, device="cuda", dtype=torch.float32) + base = torch.arange(3 * self.H, dtype=torch.float32).reshape(3, 1, self.H) + routes = torch.arange(3 * self.top_k).reshape(3, 1, self.top_k) + + def inputs(step): + hidden.copy_((((base + step * 7) % 31 - 15) / 64).bfloat16()) + cpu_ids = ((routes + step) % self.num_experts).int() + cpu_ids[step % 3].fill_(-1) + ids.copy_(cpu_ids) + weights.copy_(((routes + step) % 7 + 1).float() / (7 * self.top_k)) + + def forward(): + return [self.decode(i % self.num_layers, hidden[i], weights[i], ids[i]) for i in range(3)] + + inputs(0) + forward() + stream.synchronize() + 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() + # Eager must follow replay: otherwise it can prefill shared IO and hide stale graph inputs. + eager = forward() + stream.synchronize() + self.raise_if_unhealthy() + if any(not torch.equal(got, expected) for got, expected in zip(captured, eager)): + return "graph output differs from eager with distinct task inputs" + 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 +610,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: 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/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index edd9ddc9d..04d3b6f2d 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -545,6 +545,46 @@ 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", ["kernel-probe", "graph-data", "fp16"]) +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 == "kernel-probe": + monkeypatch.setattr(_cpu_moe, "kernel_io_probe", lambda stream: "injected kernel IO rejection") + 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.float16 if fallback == "fp16" else torch.bfloat16 + assert not ex.graph_capture_safe(dtype) + if fallback == "graph-data": + assert ex._kernel_io, "must test rejection of the full data protocol, not the kernel capability probe" + 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() + + 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..f6cd4ef7b 100644 --- a/tests/moe/test_cpu_moe_fp8_block.py +++ b/tests/moe/test_cpu_moe_fp8_block.py @@ -261,3 +261,55 @@ 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)