fix(moe): prevent stale cpu expert inputs in rocm graphs

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."""

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:

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
 

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
 

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)