fix(rocm): make CPU/Hybrid MoE graph replay safe

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

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 "