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

docs/install.md

@@ -2,7 +2,9 @@
 
 ## Requirements
 
-- Linux x86_64, NVIDIA GPU, driver r580+ (CUDA 13)
+- Linux x86_64 with either:
+  - NVIDIA GPU, driver r580+ (CUDA 13), or
+  - AMD RDNA3/RDNA4 GPU (`gfx1100`-`gfx1103`, `gfx1200`, or `gfx1201`) with ROCm 7.14
 - Python >= 3.10, with [uv](https://docs.astral.sh/uv/) recommended (plain
   `pip` + `venv` works too)
 
@@ -15,6 +17,33 @@ uv pip install "freetoken[accel]"
 
 CUDA kernels are JIT-compiled on first use, need a CUDA 13 toolkit with `nvcc` on PATH.
 
+### AMD ROCm source install (experimental)
+
+Use an official ROCm PyTorch image whose PyTorch version satisfies the project's
+`torch>=2.11,<2.12` constraint. For RDNA4, the matching ROCm 7.14 image is:
+
+```bash
+VIDEO_GID="$(getent group video | cut -d: -f3)"
+RENDER_GID="$(getent group render | cut -d: -f3)"
+docker run --rm -it \
+  --device=/dev/kfd --device=/dev/dri \
+  --group-add="$VIDEO_GID" --group-add="$RENDER_GID" --ipc=host \
+  --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \
+  -e PYTORCH_ROCM_ARCH=gfx1201 -e FREETOKEN_ROCM_ARCH=gfx1201 \
+  -v "$PWD:/workspace/FreeToken" -w /workspace/FreeToken \
+  rocm/pytorch:rocm7.14_ubuntu24.04_py3.12_pytorch_release_2.11.0 bash
+```
+
+Inside the container, preserve the ROCm-enabled PyTorch already supplied by the
+image and disable build isolation so it is also used to compile the extensions:
+
+```bash
+python -m pip install --no-build-isolation -e .
+```
+
+Set both architecture variables to `gfx1200` for RX 9060 family GPUs, or to the
+actual target reported by `rocminfo`.
+
 ## Method 2: Install from source
 
 ```bash

docs/models.md

@@ -31,6 +31,19 @@ for them; other checkpoints of the same architectures work too.
   `offload`, upgraded to `hybrid` when a cached `ft bench bw` profile
   recommends it.
 
+### 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.
+
 ## Notes
 
 - `ft checkpoint` conversion is optional — it pre-converts a checkpoint into

pyproject.toml

@@ -23,6 +23,7 @@ classifiers = [
     "Intended Audience :: Science/Research",
     "Operating System :: POSIX :: Linux",
     "Environment :: GPU :: NVIDIA CUDA",
+    "Environment :: GPU :: AMD ROCm",
     "Programming Language :: Python :: 3",
     "Programming Language :: Python :: 3.10",
     "Programming Language :: Python :: 3.11",
@@ -57,7 +58,10 @@ dependencies = [
     "torch>=2.11,<2.12",
     "tqdm>=4.66,<5",
     "transformers>=5.5,<6",
-    "triton==3.6.0; platform_system == 'Linux'",
+    # CUDA torch 2.11 resolves Triton 3.6; AMD's ROCm 7.14 image supplies its
+    # gfx1201-enabled Triton 3.7 build. Keep both supported without replacing the
+    # runtime-specific wheel selected by the PyTorch distribution.
+    "triton>=3.6,<3.8; platform_system == 'Linux'",
     "uvicorn>=0.30,<1",
 ]
 

python/freetoken/engine/engine.py

@@ -15,10 +15,18 @@
 from freetoken.moe import create_moe_backend, is_offload_moe_backend
 from freetoken.moe.expert_banks import load_expert_banks
 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
@@ -30,6 +38,52 @@
 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) -> 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()
+        or executor is None
+        or executor.graph_capture_safe
+        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; "
+        "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."""
@@ -313,6 +367,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)}")
 
@@ -333,6 +388,10 @@ def __init__(self, config: EngineConfig):
         self.cpu_moe_executor = None
         if is_offload_moe_backend(config.moe_backend):
             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)
         if hasattr(self.model, "prepare_for_runtime"):
             self.model.prepare_for_runtime()
 
@@ -683,8 +742,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,
+        )
         # gpt-oss mxfp4 carries clamped-swiglu scalars; other formats use the defaults.
         executor = CpuMoeExecutor(
             cache,
@@ -694,6 +765,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=getattr(sample, "hidden_act_alpha", 1.702),
             swiglu_limit=getattr(sample, "swiglu_limit", None),
         )
@@ -1363,7 +1437,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 "

python/freetoken/kernel/__main__.py

@@ -6,27 +6,32 @@ def generate_clangd():
     import subprocess
 
     from freetoken.kernel.utils import DEFAULT_INCLUDE
-    from freetoken.utils import init_logger
+    from freetoken.utils import get_rocm_gfx_arch, init_logger, is_rocm
     from tvm_ffi.libinfo import find_dlpack_include_path, find_include_path
 
     logger = init_logger(__name__)
     logger.info("Generating .clangd file...")
     include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE
-    status = subprocess.run(
-        args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
-        capture_output=True,
-        check=True,
-    )
-    compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
-    major, minor = compute_cap.split(".")
+
+    # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries.
+    if is_rocm():
+        arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"]
+    else:
+        try:
+            status = subprocess.run(
+                args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"],
+                capture_output=True,
+                check=True,
+            )
+            compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0]
+            major, minor = compute_cap.split(".")
+        except (subprocess.CalledProcessError, FileNotFoundError, ValueError):
+            import torch
+
+            major, minor = torch.cuda.get_device_capability()
+        arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"]
     compile_flags = ",\n    ".join(
-        [
-            "-xcuda",
-            f"--cuda-gpu-arch=sm_{major}{minor}",
-            "-std=c++20",
-            "-Wall",
-            "-Wextra",
-        ]
+        arch_flags + ["-std=c++20", "-Wall", "-Wextra"]
         + [f"-isystem{path}" for path in include_paths]
     )
     clangd_content = f"""

python/freetoken/kernel/_toolchain.py

@@ -1,4 +1,4 @@
-"""CUDA toolchain/torch consistency checks.
+"""CUDA/HIP toolchain/torch consistency checks.
 
 Standalone on purpose: setup.py and the kernel-cache build backend load this
 file by path, so it must not import the freetoken package.
@@ -16,6 +16,11 @@
 _TRUE_VALUES = {"1", "true", "yes", "on"}
 
 
+def _is_rocm() -> bool:
+    import torch
+    return getattr(torch.version, "hip", None) is not None
+
+
 def _nvcc_path() -> str | None:
     from torch.utils.cpp_extension import CUDA_HOME
 
@@ -49,6 +54,8 @@ def check_nvcc_matches_torch() -> None:
     nvcc-built binaries link libcudart.so.<nvcc major>; at runtime only the
     torch wheel's own CUDA runtime is guaranteed to be loadable.
     """
+    if _is_rocm():
+        return  # ROCm uses hipcc, not nvcc
     if os.getenv(ALLOW_MISMATCH_ENV, "").strip().lower() in _TRUE_VALUES:
         return
     torch_major = torch_cuda_major()

python/freetoken/kernel/backend.py

@@ -42,6 +42,24 @@ def is_triton_kernels_installed() -> bool:
     return _importable("triton_kernels")
 
 
+@functools.cache
+def is_rocm() -> bool:
+    """True when torch is built for ROCm (AMD GPU)."""
+    import torch
+    return getattr(torch.version, "hip", None) is not None
+
+
+@functools.cache
+def driver_hip_version() -> int | None:
+    """ROCm driver version, or None if undetermined."""
+    # TODO(ROCm): flashinfer/sgl_kernel have no ROCm builds — Triton fallback is used.
+    try:
+        from freetoken.kernel.pinned import _load_pinned_extension
+        return int(_load_pinned_extension().driver_cuda_version()) or None
+    except Exception:
+        return None
+
+
 @functools.cache
 def driver_cuda_version() -> int | None:
     """Max CUDA version the installed NVIDIA driver supports (``13000`` == CUDA 13.0),