From 4aa04e5da2b004e23b722c4d6b84d7b7d0dea5d8 Mon Sep 17 00:00:00 2001 From: Aron van Ammers Date: Tue, 15 Sep 2026 10:20:54 +0200 Subject: [PATCH] feat(windows): gfx1102 CPU-MoE serve stack for 8 GB cards Build _cpu_moe with clang-cl, Windows physical-core topology, skip missing torch.distributed, and keep ZMQ TCP ports below the ephemeral range. Validated on Framework 16 / RX 7700S / gpt-oss-20b at ~10.5 tok/s with hybrid fetch disabled. Co-authored-by: Cursor --- WINDOWS-ROCM.md | 37 +++++ pyproject.toml | 3 +- python/freetoken/engine/engine.py | 18 ++- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 21 ++- python/freetoken/moe/cpu_executor.py | 62 ++++++- python/freetoken/scheduler/config.py | 19 +-- python/freetoken/scheduler/io.py | 10 +- python/freetoken/server/api_server.py | 12 +- python/freetoken/server/args.py | 17 +- python/freetoken/utils/mp.py | 26 +++ setup.py | 153 +++++++++++++++++- tests/utils/test_zmq_addr.py | 52 ++++++ 12 files changed, 377 insertions(+), 53 deletions(-) create mode 100644 WINDOWS-ROCM.md create mode 100644 tests/utils/test_zmq_addr.py diff --git a/WINDOWS-ROCM.md b/WINDOWS-ROCM.md new file mode 100644 index 000000000..6f1d2c1e5 --- /dev/null +++ b/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 diff --git a/pyproject.toml b/pyproject.toml index 8bd653f87..9e3b663d9 100644 --- a/pyproject.toml +++ b/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] diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 88b837fad..b2ebebf25 100644 --- a/python/freetoken/engine/engine.py +++ b/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() diff --git a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp index 880e8637a..4f0723d45 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/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 #include #include @@ -29,9 +38,19 @@ #include #include +#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM) +#include +#else #include +#endif #include +#if defined(_WIN32) +#include +#else +#include +#endif + #if defined(__linux__) #include #include @@ -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 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 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); diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index b96205aa6..ffdff70fe 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/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: diff --git a/python/freetoken/scheduler/config.py b/python/freetoken/scheduler/config.py index f213f7b10..5ce004739 100644 --- a/python/freetoken/scheduler/config.py +++ b/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: diff --git a/python/freetoken/scheduler/io.py b/python/freetoken/scheduler/io.py index 37557a97f..3446eed88 100644 --- a/python/freetoken/scheduler/io.py +++ b/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] = [] diff --git a/python/freetoken/server/api_server.py b/python/freetoken/server/api_server.py index f5ef08510..7c12cac52 100644 --- a/python/freetoken/server/api_server.py +++ b/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) diff --git a/python/freetoken/server/args.py b/python/freetoken/server/args.py index d501a5eb9..324abd579 100644 --- a/python/freetoken/server/args.py +++ b/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 diff --git a/python/freetoken/utils/mp.py b/python/freetoken/utils/mp.py index 47040d976..601143495 100644 --- a/python/freetoken/utils/mp.py +++ b/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__( diff --git a/setup.py b/setup.py index e119290ca..bc1f9b125 100644 --- a/setup.py +++ b/setup.py @@ -1,13 +1,19 @@ from __future__ import annotations import importlib.util +import os +import shutil +import subprocess +import warnings from pathlib import Path from setuptools import setup -from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension +from torch.utils.cpp_extension import BuildExtension, CUDA_HOME, CppExtension, ROCM_HOME ROOT = Path(__file__).parent +KERNEL_INCLUDE = str(ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include") +IS_WINDOWS = os.name == "nt" def _check_toolchain() -> None: @@ -18,6 +24,45 @@ def _check_toolchain() -> None: module.check_nvcc_matches_torch() +def _is_rocm() -> bool: + import torch + + return getattr(torch.version, "hip", None) is not None + + +def _rocm_paths() -> tuple[list[str], list[str], str]: + candidates: list[Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(Path(os.environ["ROCM_HOME"])) + if ROCM_HOME: + candidates.append(Path(ROCM_HOME)) + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(Path(next(iter(spec.submodule_search_locations)))) + candidates.append(Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + include_dir = rocm_home / "include" + library_dir = rocm_home / "lib" + if not (include_dir / "hip" / "hip_runtime.h").exists(): + continue + if IS_WINDOWS: + if (library_dir / "amdhip64.lib").exists(): + return [str(include_dir)], [str(library_dir)], "amdhip64" + continue + if (library_dir / "libamdhip64.so").exists(): + return [str(include_dir)], [str(library_dir)], "amdhip64" + versioned = sorted(library_dir.glob("libamdhip64.so.*")) + if versioned: + return [str(include_dir)], [str(library_dir)], f":{versioned[-1].name}" + + searched = ", ".join(str(path) for path in dict.fromkeys(candidates)) + raise RuntimeError( + "A ROCm SDK with HIP headers and libamdhip64 is required to build on ROCm; " + f"searched: {searched}. Set ROCM_HOME to override." + ) + + def _cuda_runtime_paths() -> tuple[list[str], list[str]]: if CUDA_HOME is None: raise RuntimeError( @@ -31,15 +76,109 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: return [str(cuda_home / "include")], library_dirs -# patched: CUDA-only extensions are optional; skip them when no CUDA toolchain -# is present (e.g. ROCm builds use torch's own pinned-memory path instead). -import os +def _clang_cxx() -> str | None: + cxx = shutil.which(os.environ.get("CXX") or "") or "" + return cxx if "clang" in Path(cxx).stem.lower() else None -cuda_include_dirs, cuda_library_dirs = [], [] -ext_modules = [] -if os.environ.get("FREETOKEN_SKIP_CUDA_EXT") != "1" and CUDA_HOME is not None: + +def _clang_rt_builtins() -> list[str]: + cxx = _clang_cxx() + if not IS_WINDOWS or cxx is None: + return [] + resource_dir = subprocess.run( + [cxx, "-print-resource-dir"], capture_output=True, text=True, check=True + ).stdout.strip() + lib = Path(resource_dir) / "lib" / "windows" / "clang_rt.builtins-x86_64.lib" + if not lib.exists(): + raise RuntimeError( + f"{cxx} is the configured compiler but {lib} is missing; install LLVM's " + "compiler-rt component (freetoken.kernel._cpu_moe links against it)." + ) + return [str(lib)] + + +def _cpu_moe_extensions( + extra_compile: list[str], + thread_compile_args: list[str], + runtime_include_dirs: list[str], + runtime_library_dirs: list[str], + runtime_lib: str, + runtime_link_args: list[str], +) -> list[CppExtension]: + if IS_WINDOWS and _clang_cxx() is None: + warnings.warn( + "freetoken.kernel._cpu_moe is not being built: its runtime ISA dispatch " + "needs a clang driver, and CXX is unset or MSVC. Set CXX=clang-cl to build " + "it; without it --moe-backend cpu and hybrid are unavailable.", + stacklevel=2, + ) + return [] + compile_args = extra_compile + thread_compile_args + if _is_rocm(): + compile_args = compile_args + ["-D__HIP_PLATFORM_AMD__=1", "-DUSE_ROCM=1"] + return [ + CppExtension( + name="freetoken.kernel._cpu_moe", + sources=[ + "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", + ], + include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=compile_args, + extra_link_args=runtime_link_args + _clang_rt_builtins(), + ) + ] + + +IS_ROCM = _is_rocm() +if IS_WINDOWS: + extra_compile = ["/O2", "/std:c++17"] + thread_compile_args: list[str] = [] +else: + extra_compile = ["-O3", "-std:c++17"] + thread_compile_args = ["-pthread"] + +if IS_ROCM: + runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths() + runtime_link_args = [] if IS_WINDOWS else [f"-Wl,-rpath,{runtime_library_dirs[0]}"] +else: + runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() if CUDA_HOME else ([], []) + runtime_lib = "cudart" + runtime_link_args = [] + +# CUDA-only _pinned_tensor is optional; skip it when no CUDA toolchain is present +# (ROCm builds use torch's own pinned-memory path instead). +ext_modules: list[CppExtension] = [] +if os.environ.get("FREETOKEN_SKIP_CUDA_EXT") != "1" and CUDA_HOME is not None and not IS_ROCM: _check_toolchain() cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() + ext_modules.append( + CppExtension( + name="freetoken.kernel._pinned_tensor", + sources=[ + "python/freetoken/kernel/csrc/pinned_tensor.cpp", + ], + include_dirs=[KERNEL_INCLUDE, *cuda_include_dirs], + library_dirs=cuda_library_dirs, + libraries=["cudart"], + extra_compile_args=extra_compile, + ) + ) + +if IS_ROCM or CUDA_HOME is not None: + if not IS_ROCM and CUDA_HOME is not None: + _check_toolchain() + ext_modules.extend( + _cpu_moe_extensions( + extra_compile, + thread_compile_args, + runtime_include_dirs, + runtime_library_dirs, + runtime_lib, + runtime_link_args, + ) + ) setup( diff --git a/tests/utils/test_zmq_addr.py b/tests/utils/test_zmq_addr.py new file mode 100644 index 000000000..4ea26e996 --- /dev/null +++ b/tests/utils/test_zmq_addr.py @@ -0,0 +1,52 @@ +import zlib + +import pytest + +from freetoken.utils import mp +from freetoken.utils.mp import NUM_ZMQ_CHANNELS, zmq_addr + + +@pytest.fixture +def tcp(monkeypatch): + """Force the loopback-TCP fallback taken where libzmq has no ipc:// transport.""" + monkeypatch.setattr(mp, "_HAS_IPC", False) + + +@pytest.fixture +def ipc(monkeypatch): + monkeypatch.setattr(mp, "_HAS_IPC", True) + + +def test_ipc_transport_used_when_available(ipc): + assert zmq_addr(2, "-1234") == "ipc:///tmp/freetoken_2-1234" + + +def test_tcp_ports_stay_below_the_dynamic_range(tcp): + # Windows hands out ephemeral ports from 49152 up; a bind there can lose to an + # unrelated outbound connection. Sweep the whole crc32 image, not a sample. + for block in (0, mp._ZMQ_PORT_BLOCKS - 1): + base = mp._ZMQ_PORT_BASE + block * NUM_ZMQ_CHANNELS + for channel in range(NUM_ZMQ_CHANNELS): + assert 1024 < base + channel < 49152 + + +def test_tcp_channels_are_distinct_and_deterministic(tcp): + ports = {zmq_addr(channel, "-777") for channel in range(NUM_ZMQ_CHANNELS)} + assert len(ports) == NUM_ZMQ_CHANNELS + assert zmq_addr(0, "-777") == zmq_addr(0, "-777") + + +def test_tcp_address_does_not_depend_on_process_hash_salt(tcp): + block = zlib.crc32(b"-4242") % mp._ZMQ_PORT_BLOCKS + port = mp._ZMQ_PORT_BASE + block * NUM_ZMQ_CHANNELS + 3 + assert zmq_addr(3, "-4242") == f"tcp://127.0.0.1:{port}" + + +def test_distinct_instances_get_distinct_blocks(tcp): + assert zmq_addr(0, "-1") != zmq_addr(0, "-2") + + +@pytest.mark.parametrize("channel", [-1, NUM_ZMQ_CHANNELS]) +def test_channel_id_is_bounds_checked(channel): + with pytest.raises(AssertionError): + zmq_addr(channel, "-1")