feat(windows): gfx1102 CPU-MoE serve stack (hybrid, 8 GB)

WINDOWS-ROCM.md

@@ -0,0 +1,37 @@
+# Windows gfx1102 CPU-MoE bundle (for Maxritz/FreeToken-ROCm)
+
+This branch is the Windows/ROCm serve stack that ran `openai/gpt-oss-20b` at
+~10.5 tok/s on a Framework Laptop 16 (RX 7700S, gfx1102, 8 GB). It is meant
+for [Maxritz/FreeToken-ROCm](https://github.com/Maxritz/FreeToken-ROCm), not
+as a dump of FlashML PR #132.
+
+Companion tvm-ffi overlay (hipcc JIT):
+https://github.com/AronVanAmmers/tvm-ffi/tree/windows-hip-0.1.13-post3
+
+## Hardware
+
+- Machine: Framework Laptop 16
+- GPU: Radeon RX 7700S, gfx1102, 8 GB (pin HIP_VISIBLE_DEVICES=0; 780M is device 1)
+- CPU: Ryzen 9 7940HS, 8 cores / 16 threads, AVX-512-BF16
+- RAM: 96 GB installed, 8.2 GB hardware reserved (~88 GB visible to Windows)
+- OS: Windows 11
+- HIP: TheRock nightly gfx110X-all (ROCm 7.14.0a20260612)
+- Torch: 2.11.0+rocm7.14.0a20260612
+- Model: openai/gpt-oss-20b HF MXFP4, 12.84 GB
+
+## Serve command (the 10.5 tok/s path)
+
+8 GB cannot fused-hold this checkpoint. Live PCIe expert fetch was ~0.02 tok/s.
+CPU experts with fetch disabled is what served:
+
+```powershell
+$env:HIP_VISIBLE_DEVICES = "0"
+$env:TVM_FFI_ROCM_ARCH_LIST = "gfx1102"
+python -m freetoken.cli serve --model C:\models\gpt-oss-20b `
+  --moe-backend hybrid --moe-cpu-threads 8 --moe-hybrid-max-fetch 0 `
+  --moe-cache-auto --host 127.0.0.1 --port 1919
+```
+
+`--moe-hybrid-max-fetch 0` is a bypass, not a HIP gather fix.
+
+Thread: https://github.com/FlashML-org/FreeToken/issues/82

pyproject.toml

@@ -58,7 +58,8 @@ dependencies = [
     "tqdm>=4.66,<5",
     "transformers>=5.5,<6",
     "triton==3.6.0; platform_system == 'Linux'",
-    "uvicorn>=0.30,<1",
+    # >=0.36 for Config(loop=...) accepting a loop-factory import string.
+    "uvicorn>=0.36,<1",
 ]
 
 [project.urls]

python/freetoken/engine/engine.py

@@ -422,7 +422,13 @@ def __init__(self, config: EngineConfig):
             # Prefill runs on the first comma part; warm its autotune cache.
             self._warmup_prefill()
 
-    def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup:
+    def _init_communication(self, config: EngineConfig) -> torch.distributed.ProcessGroup | None:
+        # TheRock's ROCm Windows wheels are built USE_DISTRIBUTED=0, so
+        # torch.distributed is a stub. A single rank has nothing to communicate,
+        # so run without a process group; callers guard on None.
+        if config.tp_info.size == 1 and not torch.distributed.is_available():
+            logger.info_rank0("torch.distributed unavailable; running single-rank without a process group")
+            return None
         if config.tp_info.size == 1 or config.use_pynccl:
             torch.distributed.init_process_group(
                 backend="gloo",
@@ -659,9 +665,10 @@ def _sync_get_memory(self) -> Tuple[int, int]:
         torch.cuda.reset_peak_memory_stats(self.device)
         free_memory = get_free_memory(self.device)
         free_mem_tensor = torch.tensor([free_memory, -free_memory], device="cpu", dtype=torch.int64)
-        torch.distributed.all_reduce(
-            free_mem_tensor, op=torch.distributed.ReduceOp.MIN, group=self.tp_cpu_group
-        )
+        if self.tp_cpu_group is not None:
+            torch.distributed.all_reduce(
+                free_mem_tensor, op=torch.distributed.ReduceOp.MIN, group=self.tp_cpu_group
+            )
         min_free_memory = int(free_mem_tensor[0].item())
         max_free_memory = -int(free_mem_tensor[1].item())
         if max_free_memory - min_free_memory > 2 * 1024 * 1024 * 1024:
@@ -955,7 +962,8 @@ def _warmup_prefill(self) -> None:
 
     def shutdown(self) -> None:
         self.graph_runner.destroy_cuda_graphs()
-        torch.distributed.destroy_process_group()
+        if self.tp_cpu_group is not None:
+            torch.distributed.destroy_process_group()
         destroy_distributed()
 
 

python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp

@@ -15,6 +15,15 @@
 // stored bf16 to match the GPU decode path. ISA is chosen once at construction
 // (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar).
 
+#if defined(_WIN32)
+#ifndef WIN32_LEAN_AND_MEAN
+#define WIN32_LEAN_AND_MEAN 1
+#endif
+#ifndef NOMINMAX
+#define NOMINMAX 1
+#endif
+#endif
+
 #include <algorithm>
 #include <atomic>
 #include <condition_variable>
@@ -29,9 +38,19 @@
 #include <thread>
 #include <vector>
 
+#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM)
+#include <freetoken/hip_compat.cuh>
+#else
 #include <cuda_runtime_api.h>
+#endif
 #include <torch/extension.h>
 
+#if defined(_WIN32)
+#include <windows.h>
+#else
+#include <dlfcn.h>
+#endif
+
 #if defined(__linux__)
 #include <pthread.h>
 #include <sched.h>
@@ -568,13 +587,11 @@ float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float
 // probed functionally at startup (memops_probe); anything unsupported (Windows WDDM,
 // vGPU, old drivers) falls back to the cudaLaunchHostFunc path.
 #if defined(_WIN32)
-#include <windows.h>
 static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); }
 static void* cumemop_dlsym(void* h, const char* n) {
   return (void*)::GetProcAddress((HMODULE)h, n);
 }
 #else
-#include <dlfcn.h>
 static void* cumemop_dlopen() {
   void* h = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL);
   if (h == nullptr) h = dlopen("libcuda.so", RTLD_LAZY | RTLD_LOCAL);

python/freetoken/moe/cpu_executor.py

@@ -88,14 +88,74 @@ def compiled_extension_supports(activation: str) -> bool:
     return _ACT_IDS[activation] <= getattr(_cpu_moe, "max_generic_act_id", lambda: 2)()
 
 
+def _windows_physical_core_cpus() -> list[int]:
+    """One logical CPU per physical core on Windows.
+
+    ``GetLogicalProcessorInformationEx(RelationProcessorCore)`` returns one record
+    per physical core, each carrying a group affinity mask over that core's SMT
+    siblings; the lowest set bit is the core's representative CPU. Processor
+    groups are 64 CPUs wide, which is how the ids flatten into the same numbering
+    the Linux branch uses. The whole machine is reported: Windows affinity is
+    per-group and ``_cpu_moe`` does not pin threads there anyway
+    (``CPU_MOE_HAS_AFFINITY`` is Linux-only).
+    """
+    import ctypes
+
+    class GroupAffinity(ctypes.Structure):
+        _fields_ = [
+            ("Mask", ctypes.c_size_t),
+            ("Group", ctypes.c_uint16),
+            ("Reserved", ctypes.c_uint16 * 3),
+        ]
+
+    class ProcessorRelationship(ctypes.Structure):
+        _fields_ = [
+            ("Flags", ctypes.c_uint8),
+            ("EfficiencyClass", ctypes.c_uint8),
+            ("Reserved", ctypes.c_uint8 * 20),
+            ("GroupCount", ctypes.c_uint16),
+            ("GroupMask", GroupAffinity * 1),
+        ]
+
+    class ProcessorInfoEx(ctypes.Structure):
+        _fields_ = [
+            ("Relationship", ctypes.c_uint32),
+            ("Size", ctypes.c_uint32),
+            ("Processor", ProcessorRelationship),
+        ]
+
+    relation_processor_core = 0
+    kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
+    size = ctypes.c_uint32(0)
+    kernel32.GetLogicalProcessorInformationEx(relation_processor_core, None, ctypes.byref(size))
+    buffer = (ctypes.c_char * size.value)()
+    if not kernel32.GetLogicalProcessorInformationEx(
+        relation_processor_core, buffer, ctypes.byref(size)
+    ):
+        raise ctypes.WinError(ctypes.get_last_error())
+
+    cpus: list[int] = []
+    offset = 0
+    while offset < size.value:
+        record = ProcessorInfoEx.from_buffer(buffer, offset)
+        affinity = record.Processor.GroupMask[0]
+        if affinity.Mask:
+            lowest = (affinity.Mask & -affinity.Mask).bit_length() - 1
+            cpus.append(affinity.Group * 64 + lowest)
+        offset += record.Size
+    return sorted(cpus)
+
+
 def physical_core_cpus() -> list[int]:
     """One logical CPU per physical core, restricted to this process's affinity.
 
     MoE decode is memory-bandwidth-bound, so SMT siblings only contend for the
     same core's load ports without adding bandwidth. Picking one logical CPU per
     physical core (and pinning to it) gives the best, most stable bandwidth.
-    Falls back to the full affinity set when sysfs topology is unavailable.
+    Falls back to the full affinity set when the topology is unavailable.
     """
+    if os.name == "nt":
+        return _windows_physical_core_cpus()
     try:
         allowed = sorted(os.sched_getaffinity(0))
     except AttributeError:

python/freetoken/scheduler/config.py

@@ -3,18 +3,7 @@
 from dataclasses import dataclass, field
 
 from freetoken.engine import EngineConfig
-
-def _zmq_addr(name: str) -> str:
-    """patched: ipc:// is unsupported on Windows; use localhost TCP there."""
-    import hashlib
-    import os
-    import sys
-
-    if sys.platform == "win32":
-        # patched: name-only hash - all workers must derive the SAME port
-        port = 29876 + int(hashlib.sha1(name.encode()).hexdigest()[:6], 16) % 20000
-        return f"tcp://127.0.0.1:{port}"
-    return f"ipc:///tmp/{name}"
+from freetoken.utils.mp import zmq_addr
 
 
 def _get_pid_suffix() -> str:
@@ -36,15 +25,15 @@ class SchedulerConfig(EngineConfig):
 
     @property
     def zmq_backend_addr(self) -> str:
-        return _zmq_addr("freetoken_0")
+        return zmq_addr(0, self._unique_suffix)
 
     @property
     def zmq_detokenizer_addr(self) -> str:
-        return _zmq_addr("freetoken_1")
+        return zmq_addr(1, self._unique_suffix)
 
     @property
     def zmq_scheduler_broadcast_addr(self) -> str:
-        return _zmq_addr("freetoken_2")
+        return zmq_addr(2, self._unique_suffix)
 
     @property
     def max_forward_len(self) -> int:

python/freetoken/scheduler/io.py

@@ -24,7 +24,9 @@ class SchedulerIOMixin:
         sync_all_ranks: Function to synchronize all ranks on CPU side.
     """
 
-    def __init__(self, config: SchedulerConfig, tp_cpu_group: torch.distributed.ProcessGroup):
+    def __init__(
+        self, config: SchedulerConfig, tp_cpu_group: torch.distributed.ProcessGroup | None
+    ):
         tp_info = config.tp_info
         self.tp_cpu_group: Final = tp_cpu_group
         if config.offline_mode:
@@ -74,7 +76,11 @@ def offline_send_result(self, reply: List[BaseTokenizerMsg]) -> None:
         raise NotImplementedError("should be implemented")
 
     def sync_all_ranks(self) -> None:
-        self.tp_cpu_group.barrier().wait()
+        # None when torch.distributed is unavailable; only ever single-rank then,
+        # so there is nothing to synchronize. The multi-rank broadcasts below are
+        # unreachable in that case and need no guard.
+        if self.tp_cpu_group is not None:
+            self.tp_cpu_group.barrier().wait()
 
     def _recv_msg_single_rank(self, blocking: bool = False) -> List[BaseBackendMsg]:
         pending_msgs: List[BaseBackendMsg] = []

python/freetoken/server/api_server.py

@@ -1,10 +1,6 @@
 from __future__ import annotations
 
 import asyncio
-import sys
-if sys.platform == 'win32':
-    # patched: zmq.asyncio requires a Selector loop; Windows defaults to Proactor
-    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())
 import contextlib
 import json
 import os
@@ -61,6 +57,10 @@
 # shutdown is treated as expected — no ERROR log, no "failed" latch. See run_backend_supervisor.
 _SHUTTING_DOWN = threading.Event()
 BACKEND_DEATH_EXIT_GRACE_S = 10.0
+# zmq.asyncio needs add_reader. Windows Proactor does not implement it; uvicorn's
+# "auto" loop is a Proactor there, so name the selector loop as a loop_factory
+# import string (uvicorn >= 0.36) instead of mutating the process-wide policy.
+UVICORN_LOOP = "asyncio:SelectorEventLoop" if os.name == "nt" else "auto"
 
 
 def get_global_state() -> FrontendManager:
@@ -899,7 +899,7 @@ def _serve_and_run_shell(host: str, port: int) -> None:
     netloc = f"[{host}]:{port}" if ":" in host else f"{host}:{port}"
     origin = resolve_server_url(f"http://{netloc}").origin
 
-    server = uvicorn.Server(uvicorn.Config(app, host=host, port=port, access_log=False))
+    server = uvicorn.Server(uvicorn.Config(app, host=host, port=port, access_log=False, loop=UVICORN_LOOP))
     thread = threading.Thread(target=server.run, name="freetoken-uvicorn", daemon=True)
     thread.start()
     _install_shell_stop_handlers()
@@ -1038,4 +1038,4 @@ def _on_meta(meta: dict) -> None:
         _serve_and_run_shell(host, port)
         return
     # uvicorn stays on the main thread (signal handling unchanged); ^C reaches the worker group.
-    uvicorn.run(app, host=host, port=port)
+    uvicorn.run(app, host=host, port=port, loop=UVICORN_LOOP)

python/freetoken/server/args.py

@@ -9,18 +9,7 @@
 from freetoken.distributed import DistributedInfo
 from freetoken.scheduler import SchedulerConfig
 from freetoken.utils import init_logger
-
-def _zmq_addr(name: str) -> str:
-    """patched: ipc:// is unsupported on Windows; use localhost TCP there."""
-    import hashlib
-    import os
-    import sys
-
-    if sys.platform == "win32":
-        # patched: name-only hash - all workers must derive the SAME port
-        port = 29876 + int(hashlib.sha1(name.encode()).hexdigest()[:6], 16) % 20000
-        return f"tcp://127.0.0.1:{port}"
-    return f"ipc:///tmp/{name}"
+from freetoken.utils.mp import zmq_addr
 
 
 @dataclass(frozen=True)
@@ -58,13 +47,13 @@ def share_tokenizer(self) -> bool:
 
     @property
     def zmq_frontend_addr(self) -> str:
-        return _zmq_addr("freetoken_3")
+        return zmq_addr(3, self._unique_suffix)
 
     @property
     def zmq_tokenizer_addr(self) -> str:
         if self.share_tokenizer:
             return self.zmq_detokenizer_addr
-        result = _zmq_addr("freetoken_4")
+        result = zmq_addr(4, self._unique_suffix)
         assert result != self.zmq_detokenizer_addr
         return result
 

python/freetoken/utils/mp.py

@@ -2,6 +2,7 @@
 
 from typing import Callable, Dict, Generic, TypeVar
 import sys
+import zlib
 
 if sys.platform == 'win32':
     # patched: zmq.asyncio needs a Selector loop; win32 defaults to Proactor
@@ -14,6 +15,31 @@
 
 T = TypeVar("T")
 
+# Number of internal channels an instance uses (see SchedulerConfig and ServerArgs
+# for the ids). On the TCP fallback below it is also the port stride between two
+# instances, so every instance gets a contiguous block of its own.
+NUM_ZMQ_CHANNELS = 5
+# Highest port handed out is _ZMQ_PORT_BASE + _ZMQ_PORT_BLOCKS * NUM_ZMQ_CHANNELS - 1
+# = 44999, below 49152 where Windows' default dynamic-port range begins.
+_ZMQ_PORT_BASE = 20000
+_ZMQ_PORT_BLOCKS = 5000
+_HAS_IPC = zmq.has("ipc")
+
+
+def zmq_addr(channel: int, suffix: str) -> str:
+    """ZeroMQ endpoint for one of an instance's internal channels.
+
+    Falls back to a loopback TCP port where libzmq was built without the
+    ``ipc://`` transport (every Windows build). The port is derived from the
+    per-instance ``suffix`` with crc32 rather than ``hash``: hash() is salted per
+    process, and every process of the instance has to arrive at the same address.
+    """
+    assert 0 <= channel < NUM_ZMQ_CHANNELS, channel
+    if _HAS_IPC:
+        return f"ipc:///tmp/freetoken_{channel}{suffix}"
+    block = zlib.crc32(suffix.encode()) % _ZMQ_PORT_BLOCKS
+    return f"tcp://127.0.0.1:{_ZMQ_PORT_BASE + block * NUM_ZMQ_CHANNELS + channel}"
+
 
 class ZmqPushQueue(Generic[T]):
     def __init__(