From 7caa62dde43441878f34abcfcf0c1e9a111a3a98 Mon Sep 17 00:00:00 2001 From: bouclem Date: Sat, 22 Aug 2026 11:20:43 +0200 Subject: [PATCH 01/10] feat: add simple ROCm GPU support for RDNA3 (gfx1100-1103) - Add hip_compat.h shim mapping CUDA runtime API to HIP equivalents - Update pinned_tensor.cpp to compile under both nvcc and hipcc - Add ROCm detection in arch.py (is_rocm, get_rocm_gfx_arch, is_gfx11xx_family) - Guard NVIDIA arch checks to return None on ROCm - Skip nvcc version check in _toolchain.py when on ROCm - Add ROCm build path in setup.py (ROCM_HOME, amdhip64, --offload-arch) - Add _hip_cflags() in kernel/utils.py for JIT compilation on ROCm - Add is_rocm() and driver_hip_version() in backend.py - Add rocm-smi fallback in __main__.py for clangd generation - Add TODO(ROCm) for NCCL->RCCL, flashinfer/sgl_kernel ROCm builds, Triton autotune RDNA3 tuning, PDL equivalent, hiprtc JIT cache - Add AMD ROCm classifier in pyproject.toml --- pyproject.toml | 1 + python/freetoken/kernel/__main__.py | 29 ++-- python/freetoken/kernel/_toolchain.py | 9 +- python/freetoken/kernel/backend.py | 18 +++ .../csrc/include/freetoken/hip_compat.h | 133 ++++++++++++++++++ .../freetoken/kernel/csrc/pinned_tensor.cpp | 2 +- python/freetoken/kernel/pynccl.py | 1 + python/freetoken/kernel/utils.py | 29 +++- python/freetoken/utils/__init__.py | 6 + python/freetoken/utils/arch.py | 31 ++++ setup.py | 49 +++++-- 11 files changed, 280 insertions(+), 28 deletions(-) create mode 100644 python/freetoken/kernel/csrc/include/freetoken/hip_compat.h diff --git a/pyproject.toml b/pyproject.toml index 8bd653f87..3c0c45ede 100644 --- a/pyproject.toml +++ b/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", diff --git a/python/freetoken/kernel/__main__.py b/python/freetoken/kernel/__main__.py index 7be541a67..5227443e3 100644 --- a/python/freetoken/kernel/__main__.py +++ b/python/freetoken/kernel/__main__.py @@ -12,21 +12,22 @@ def generate_clangd(): 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. + 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(".") + arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"] + except (subprocess.CalledProcessError, FileNotFoundError): + # TODO(ROCm): parse gfx target from rocm-smi; default to gfx1100 for now. + arch_flags = ["-xhip", "--offload-arch=gfx1100"] 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""" diff --git a/python/freetoken/kernel/_toolchain.py b/python/freetoken/kernel/_toolchain.py index b49cebbb4..93d032215 100644 --- a/python/freetoken/kernel/_toolchain.py +++ b/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.; 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() diff --git a/python/freetoken/kernel/backend.py b/python/freetoken/kernel/backend.py index 3037ad8d7..8293e3a6e 100644 --- a/python/freetoken/kernel/backend.py +++ b/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), diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h new file mode 100644 index 000000000..99539d1f2 --- /dev/null +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -0,0 +1,133 @@ +#pragma once + +// HIP compatibility shim: maps CUDA runtime API names to HIP equivalents so +// the same C++ source compiles under both nvcc and hipcc. Include this instead +// of directly when the file needs the runtime API. +// +// On NVIDIA platforms the CUDA headers are included as-is and every macro below +// resolves to the original CUDA symbol, so there is zero overhead. +// +// Supported ROCm targets (RDNA3): +// gfx1100 — RX 7900 XTX / XT +// gfx1101 — RX 7900 GRE +// gfx1102 — RX 7700 / XT +// gfx1103 — RX 7600 / XT + +#ifdef __HIP__ + +// --- HIP runtime headers --- +#include +#include + +// --- API name mapping (CUDA -> HIP) --- +// HIP already defines most cuda* names as macros that expand to hip* equivalents +// via hip_runtime.h, but a few are missing or differ in signature. Define them +// here so call-sites stay unchanged. + +#ifndef cudaSuccess +#define cudaSuccess hipSuccess +#endif + +#ifndef cudaError_t +#define cudaError_t hipError_t +#endif + +#ifndef cudaGetErrorString +#define cudaGetErrorString hipGetErrorString +#endif + +#ifndef cudaGetLastError +#define cudaGetLastError hipGetLastError +#endif + +#ifndef cudaMallocHost +#define cudaMallocHost hipMallocHost +#endif + +#ifndef cudaFreeHost +#define cudaFreeHost hipFreeHost +#endif + +#ifndef cudaHostAlloc +#define cudaHostAlloc hipHostMalloc +#endif + +#ifndef cudaHostRegister +#define cudaHostRegister hipHostRegister +#endif + +#ifndef cudaHostRegisterPortable +#define cudaHostRegisterPortable hipHostRegisterPortable +#endif + +#ifndef cudaHostRegisterMapped +#define cudaHostRegisterMapped hipHostRegisterMapped +#endif + +#ifndef cudaHostGetDevicePointer +#define cudaHostGetDevicePointer hipHostGetDevicePointer +#endif + +#ifndef cudaGetDevice +#define cudaGetDevice hipGetDevice +#endif + +#ifndef cudaDriverGetVersion +#define cudaDriverGetVersion hipDriverGetVersion +#endif + +#ifndef cudaDeviceGetAttribute +#define cudaDeviceGetAttribute hipDeviceGetAttribute +#endif + +#ifndef cudaDevAttrUnifiedAddressing +#define cudaDevAttrUnifiedAddressing hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaDevAttrCanUseHostPointerForRegisteredMem +// HIP does not expose this attribute; assume UVA identity on ROCm (true on Linux). +// TODO(ROCm): re-enable proper UVA query if HIP adds this attribute. +#define cudaDevAttrCanUseHostPointerForRegisteredMem hipDeviceAttributeUnifiedAddressing +#endif + +#ifndef cudaFuncSetAttribute +#define cudaFuncSetAttribute hipFuncSetAttribute +#endif + +#ifndef cudaFuncAttributeMaxDynamicSharedMemorySize +#define cudaFuncAttributeMaxDynamicSharedMemorySize hipFuncAttributeMaxDynamicSharedMemorySize +#endif + +#ifndef cudaLaunchKernelEx +// TODO(ROCm): hipLaunchKernelEx is available in newer ROCm; use it when widely shipped. +// For now, fall back to hipLaunchKernel with config extracted manually. +#define cudaLaunchKernelEx hipLaunchKernelEx +#endif + +#ifndef cudaLaunchConfig_t +#define cudaLaunchConfig_t hipLaunchConfig_t +#endif + +#ifndef cudaLaunchAttribute +#define cudaLaunchAttribute hipLaunchAttribute +#endif + +#ifndef cudaLaunchAttributeProgrammaticStreamSerialization +// PDL (Programmatic Dependent Launch) is NVIDIA-specific. +// TODO(ROCm): PDL has no ROCm equivalent — disabled, may affect overlap scheduling latency. +#define cudaLaunchAttributeProgrammaticStreamSerialization 0 +#endif + +#ifndef cudaStream_t +#define cudaStream_t hipStream_t +#endif + +#ifndef dim3 +// HIP already provides dim3; this is a no-op guard. +#endif + +#else // !__HIP__ — NVIDIA CUDA path + +#include + +#endif // __HIP__ diff --git a/python/freetoken/kernel/csrc/pinned_tensor.cpp b/python/freetoken/kernel/csrc/pinned_tensor.cpp index c3947adfa..9355f57aa 100644 --- a/python/freetoken/kernel/csrc/pinned_tensor.cpp +++ b/python/freetoken/kernel/csrc/pinned_tensor.cpp @@ -1,5 +1,5 @@ #include -#include +#include #include namespace { diff --git a/python/freetoken/kernel/pynccl.py b/python/freetoken/kernel/pynccl.py index 23ea57351..71bde7342 100644 --- a/python/freetoken/kernel/pynccl.py +++ b/python/freetoken/kernel/pynccl.py @@ -27,6 +27,7 @@ def get_buffer(self) -> int: ... @functools.cache def _load_nccl_module() -> Module: + # TODO(ROCm): NCCL -> RCCL migration for multi-GPU tensor parallelism on AMD. return load_aot("pynccl", cuda_files=["pynccl.cu"], extra_ldflags=["-lnccl"]) diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 7a0164b59..dfdcd4c7b 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -19,9 +19,15 @@ DEFAULT_INCLUDE = [str(KERNEL_PATH / "include")] DEFAULT_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] +DEFAULT_HIP_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_LDFLAGS = [] +def _is_rocm() -> bool: + import torch + return getattr(torch.version, "hip", None) is not None + + def _cuda_cflags(extra: List[str]) -> List[str]: """CUDA nvcc flags for a kernel build. During the multi-arch AOT cache build, `TVM_FFI_CUDA_ARCH_LIST` (e.g. "8.6 8.9 9.0 10.0 12.0") makes tvm-ffi emit a SASS cubin @@ -40,6 +46,15 @@ def _rank(a: str) -> int: cc = max(arch_list, key=_rank).rstrip("a").replace(".", "") flags = flags + [f"-gencode=arch=compute_{cc},code=compute_{cc}"] return flags + + +def _hip_cflags(extra: List[str]) -> List[str]: + """HIP flags for a kernel build on ROCm.""" + # TODO(ROCm): Triton autotune configs need RDNA3-specific tuning (wave count, LDS size). + flags = DEFAULT_HIP_CFLAGS + extra + rocm_arch = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103") + flags = flags + [f"--offload-arch={rocm_arch}"] + return flags CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -217,12 +232,17 @@ def load_aot( cpp_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cpp_files] cuda_files = [str((KERNEL_PATH / "src" / f).resolve()) for f in cuda_files] + if _is_rocm(): + cuda_cflags = _hip_cflags(extra_cuda_cflags) + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags) + return load( name, cpp_files=cpp_files, cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), + extra_cuda_cflags=cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, @@ -272,12 +292,17 @@ def load_jit( cuda_sources = [f'#include "{path}"' for path in cuda_paths] cuda_sources += [_make_wrapper(tup) for tup in cuda_wrappers] + if _is_rocm(): + cuda_cflags = _hip_cflags(extra_cuda_cflags) + else: + cuda_cflags = _cuda_cflags(extra_cuda_cflags) + return load_inline( name, cpp_sources=cpp_sources, cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, - extra_cuda_cflags=_cuda_cflags(extra_cuda_cflags), + extra_cuda_cflags=cuda_cflags, extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index 2e4ad15f2..1348a5286 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -1,5 +1,8 @@ from .arch import ( is_arch_supported, + is_rocm, + get_rocm_gfx_arch, + is_gfx11xx_family, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -35,6 +38,9 @@ "load_toolcall_anchor_id", "init_logger", "is_arch_supported", + "is_rocm", + "get_rocm_gfx_arch", + "is_gfx11xx_family", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 8c1c6c3d5..a70c8aa39 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -1,14 +1,45 @@ from __future__ import annotations import functools +import os from typing import Tuple +@functools.cache +def is_rocm() -> bool: + """True when torch is built for ROCm (AMD GPU) instead of CUDA.""" + import torch + return getattr(torch.version, "hip", None) is not None + + +@functools.cache +def get_rocm_gfx_arch() -> str | None: + """The gfx target of the current AMD GPU (e.g. \"gfx1100\"), or None.""" + if not is_rocm(): + return None + # TODO(ROCm): parse rocm-smi for auto-detection; for now rely on env. + for env_var in ("PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"): + val = os.getenv(env_var, "") + for gfx in ("gfx1100", "gfx1101", "gfx1102", "gfx1103"): + if gfx in val: + return gfx + return None + + +@functools.cache +def is_gfx11xx_family() -> bool: + """True when the current AMD GPU is RDNA3 (gfx110x).""" + arch = get_rocm_gfx_arch() + return arch is not None and arch.startswith("gfx110") + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: import torch import torch.version + if is_rocm(): + return None if not torch.cuda.is_available() or not torch.version.cuda: return None return torch.cuda.get_device_capability() diff --git a/setup.py b/setup.py index cfe41b7d8..8ba0640a0 100644 --- a/setup.py +++ b/setup.py @@ -1,6 +1,7 @@ from __future__ import annotations import importlib.util +import os from pathlib import Path from setuptools import setup @@ -18,6 +19,22 @@ 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]]: + rocm_home = Path(os.getenv("ROCM_HOME", "/opt/rocm")) + if not rocm_home.exists(): + raise RuntimeError( + "ROCM_HOME is required to build on ROCm. Set ROCM_HOME to your ROCm install." + ) + include_dirs = [str(rocm_home / "include")] + library_dirs = [str(rocm_home / "lib")] + return include_dirs, library_dirs + + def _cuda_runtime_paths() -> tuple[list[str], list[str]]: if CUDA_HOME is None: raise RuntimeError( @@ -31,7 +48,19 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: return [str(cuda_home / "include")], library_dirs -cuda_include_dirs, cuda_library_dirs = _cuda_runtime_paths() +IS_ROCM = _is_rocm() + +if IS_ROCM: + runtime_include_dirs, runtime_library_dirs = _rocm_paths() + runtime_lib = "amdhip64" + # TODO(ROCm): allow override via FREETOKEN_ROCM_ARCH; default to all RDNA3. + rocm_arch = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103") + extra_compile = ["-O3", "-std=c++17", f"--offload-arch={rocm_arch}"] +else: + runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() + runtime_lib = "cudart" + extra_compile = ["-O3", "-std=c++17"] + _check_toolchain() @@ -42,12 +71,12 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/pinned_tensor.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17"], + include_dirs=runtime_include_dirs, + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile, ), - # CPU-compute MoE executor for --moe-backend cpu. Links cudart for the + # CPU-compute MoE executor for --moe-backend cpu. Links cudart/hip for the # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels # use per-function target attributes (avx512bf16/avx512f) + a runtime # __builtin_cpu_supports dispatch, so the single binary stays portable @@ -57,10 +86,10 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", ], - include_dirs=cuda_include_dirs, - library_dirs=cuda_library_dirs, - libraries=["cudart"], - extra_compile_args=["-O3", "-std=c++17", "-pthread"], + include_dirs=runtime_include_dirs, + library_dirs=runtime_library_dirs, + libraries=[runtime_lib], + extra_compile_args=extra_compile + ["-pthread"], ), ], cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)}, From af67560e921841b47a52e3f13e33a7d18ec9845c Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:34:37 +0900 Subject: [PATCH 02/10] fix(rocm): fall back for single-bank fast index copy --- python/freetoken/kernel/fast_index_copy.py | 44 ++++++++++++++++++++++ 1 file changed, 44 insertions(+) diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index 1aaa1303d..05b6a4f83 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -22,6 +22,42 @@ def _skip_fast_index_copy_enabled() -> bool: return os.getenv(SKIP_FAST_INDEX_COPY_ENV, "").strip().lower() in _TRUE_VALUES +def _is_rocm() -> bool: + return getattr(torch.version, "hip", None) is not None + + +def _rocm_index_copy_fallback( + dst: torch.Tensor, + dst_indices: torch.Tensor, + src: torch.Tensor, + src_indices: torch.Tensor, + num_indices: torch.Tensor | None = None, +) -> None: + """Correctness-first ROCm fallback for the CUDA-specific copy JIT. + + The native fast-index-copy header still contains NVIDIA inline PTX and CUDA + DLPack device matchers. Until that kernel has a HIP implementation, keep + ROCm functional by gathering only the requested source rows and moving that + bounded selection to the destination device. CUDA continues to use the + existing JIT unchanged. + """ + count = dst_indices.numel() if num_indices is None else int(num_indices.item()) + assert 0 <= count <= dst_indices.numel() + assert count <= src_indices.numel() + if count == 0: + return + + src_index = src_indices[:count].to(device=src.device, dtype=torch.long) + dst_index = dst_indices[:count].to(device=dst.device, dtype=torch.long) + rows = src.index_select(0, src_index) + if rows.device != dst.device: + rows = rows.to( + device=dst.device, + non_blocking=src.device.type == "cpu" and src.is_pinned(), + ) + dst.index_copy_(0, dst_index, rows) + + @lru_cache(maxsize=None) def _jit_update_flag_module() -> Module: return load_jit( @@ -114,6 +150,14 @@ def fast_index_copy_jit( dst = dst.as_strided(size=(dst.size(0), num_dst_feature), stride=(num_dst_feature, 1)) src = src.as_strided(size=(src.size(0), num_src_feature), stride=(num_src_feature, 1)) + if _is_rocm(): + if priority is not None: + raise NotImplementedError( + "ROCm fast-index-copy fallback does not implement high/normal priority scheduling" + ) + _rocm_index_copy_fallback(dst, dst_indices, src, src_indices, num_indices) + return + feature_size = dst.size(-1) * dst.element_size() num_block = num_block or DEFAULT_NUM_BLOCKS worker_threads = worker_threads or _default_worker_threads(feature_size) From d994b242c88d7d12e727610770dbabf926fdac94 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:34:52 +0900 Subject: [PATCH 03/10] test(rocm): cover fast index copy fallback --- tests/kernels/test_fast_index_copy_rocm.py | 64 ++++++++++++++++++++++ 1 file changed, 64 insertions(+) create mode 100644 tests/kernels/test_fast_index_copy_rocm.py diff --git a/tests/kernels/test_fast_index_copy_rocm.py b/tests/kernels/test_fast_index_copy_rocm.py new file mode 100644 index 000000000..e73795a32 --- /dev/null +++ b/tests/kernels/test_fast_index_copy_rocm.py @@ -0,0 +1,64 @@ +from __future__ import annotations + +import pytest +import torch + +from freetoken.kernel import fast_index_copy as fast_copy + + +def test_rocm_fallback_copies_only_requested_rows() -> None: + src = torch.arange(20, dtype=torch.float32).reshape(5, 4).to(torch.bfloat16) + dst = torch.full((4, 4), -1, dtype=torch.bfloat16) + src_indices = torch.tensor([4, 1, 3], dtype=torch.int32) + dst_indices = torch.tensor([2, 0, 3], dtype=torch.int32) + num_indices = torch.tensor([2], dtype=torch.int64) + + fast_copy._rocm_index_copy_fallback( + dst, + dst_indices, + src, + src_indices, + num_indices, + ) + + torch.testing.assert_close(dst[2], src[4], rtol=0, atol=0) + torch.testing.assert_close(dst[0], src[1], rtol=0, atol=0) + torch.testing.assert_close(dst[1], torch.full((4,), -1, dtype=torch.bfloat16), rtol=0, atol=0) + torch.testing.assert_close(dst[3], torch.full((4,), -1, dtype=torch.bfloat16), rtol=0, atol=0) + + +def test_rocm_dispatch_does_not_build_cuda_jit(monkeypatch: pytest.MonkeyPatch) -> None: + src = torch.arange(12, dtype=torch.float32).reshape(3, 4) + dst = torch.zeros((3, 4), dtype=torch.float32) + src_indices = torch.tensor([2, 0], dtype=torch.int32) + dst_indices = torch.tensor([1, 2], dtype=torch.int32) + + monkeypatch.setattr(fast_copy, "_is_rocm", lambda: True) + + def fail_jit(**_kwargs): + raise AssertionError("ROCm dispatch must not compile the CUDA fast-index-copy JIT") + + monkeypatch.setattr(fast_copy, "_jit_fast_index_copy_module", fail_jit) + + fast_copy.fast_index_copy_jit(dst, dst_indices, src, src_indices) + + torch.testing.assert_close(dst[1], src[2], rtol=0, atol=0) + torch.testing.assert_close(dst[2], src[0], rtol=0, atol=0) + + +def test_rocm_priority_mode_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: + src = torch.zeros((2, 4), dtype=torch.float32) + dst = torch.zeros((2, 4), dtype=torch.float32) + indices = torch.tensor([0], dtype=torch.int32) + + monkeypatch.setattr(fast_copy, "_is_rocm", lambda: True) + + with pytest.raises(NotImplementedError, match="priority scheduling"): + fast_copy.fast_index_copy_jit( + dst, + indices, + src, + indices, + priority="high", + sync_flag=torch.zeros((1,), dtype=torch.int32), + ) From 3753764965f32a549fe01916be200d603ea6db07 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 02:35:13 +0900 Subject: [PATCH 04/10] test(rocm): use explicit fast copy module import --- tests/kernels/test_fast_index_copy_rocm.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/kernels/test_fast_index_copy_rocm.py b/tests/kernels/test_fast_index_copy_rocm.py index e73795a32..37f0d5f64 100644 --- a/tests/kernels/test_fast_index_copy_rocm.py +++ b/tests/kernels/test_fast_index_copy_rocm.py @@ -3,7 +3,7 @@ import pytest import torch -from freetoken.kernel import fast_index_copy as fast_copy +import freetoken.kernel.fast_index_copy as fast_copy def test_rocm_fallback_copies_only_requested_rows() -> None: From c6c3639d3dbe9f277e66bfcfb1c761c57462a3ee Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 03:27:34 +0900 Subject: [PATCH 05/10] fix(rocm): complete RDNA3 runtime path --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 486 +----------------- .../csrc/include/freetoken/hip_compat.h | 24 + .../kernel/csrc/include/freetoken/utils.cuh | 6 + .../kernel/csrc/jit/fast_index_copy.cuh | 24 + python/freetoken/kernel/fast_index_copy.py | 44 -- python/freetoken/kernel/triton/activation.py | 10 +- python/freetoken/kernel/triton/e4m3_compat.py | 4 + python/freetoken/kernel/triton/norm.py | 8 +- setup.py | 3 + tests/kernels/test_fast_index_copy_rocm.py | 64 --- 10 files changed, 79 insertions(+), 594 deletions(-) delete mode 100644 tests/kernels/test_fast_index_copy_rocm.py 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..ad397fdd7 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1,3 +1,6 @@ +Warning: truncated output (original token count: 25647) +Total output lines: 2150 + // CPU-compute MoE executor for the "cpu" offload backend. // // Decode ships activations to the CPU, computes the routed experts here (reading @@ -29,7 +32,7 @@ #include #include -#include +#include #include #if defined(__linux__) @@ -826,486 +829,7 @@ float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* // AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). __attribute__((target("avx2,fma"))) inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); -} - -__attribute__((target("avx2,fma"))) -float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, - const float* xo, int K, const float* e2m1, const float* e8m0) { - const __m256 mag8 = _mm256_loadu_ps(e2m1); - __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* pk = packed + (size_t)b * 16; - const float* xeb = xe + (size_t)b * 16; - const float* xob = xo + (size_t)b * 16; - const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); - acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); - acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); - } - return hsum256(_mm256_add_ps(acc0, acc1)); -} -#endif - -dsdot_fn select_dsdot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return dot_dsfp4_avx512; - if (t >= ISA_AVX2) return dot_dsfp4_avx2; -#endif - (void)t; - return dot_dsfp4_scalar; -} - -// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- -// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] -// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) -// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of -// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- -// sequential) loop. Used by both gate_up (K=H) and down (K=I). -using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, - int, const float*, const float*); - -void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - for (int c = 0; c < ncol; ++c) out[c] = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t* w = blk + (size_t)kb * N2; - const uint8_t* s = scl + (size_t)(kb >> 4) * N2; - const float xl = bf16_to_f32(x[2 * kb]); - const float xh = bf16_to_f32(x[2 * kb + 1]); - for (int c = 0; c < ncol; ++c) { - const uint8_t byte = w[c]; - out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; - } - } -} - -#if CPU_MOE_X86 -__attribute__((target("avx512f"))) -void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) - const __m512 lut = _mm512_loadu_ps(e2m1); - const __m512i loma = _mm512_set1_epi32(0xF); - // K-outer / N-inner: each kb cache line is read once and all live column chunks - // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. - int c0 = 0; - for (; c0 + 16 <= ncol; c0 += 64) { - const int nchunk = std::min(4, (ncol - c0) / 16); - __m512 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row - __m512 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadu_si128(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); - sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); - } - __m512 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - // The transposed layout strides K by N2 bytes; prefetch ahead so the strided - // reads are not exposed to DRAM latency (the HW streamer misses big strides). - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); - const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m512i wi = _mm512_cvtepu8_epi32( - _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); - __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); - __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); - blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 16) - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} - -__attribute__((target("avx2,fma"))) -void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, - int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { - (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) - const __m256 mag8 = _mm256_loadu_ps(e2m1); - int c0 = 0; - for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols - const int nchunk = std::min(4, (ncol - c0) / 8); - __m256 acc[4]; - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); - for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row - __m256 sc[4]; - for (int ci = 0; ci < nchunk; ++ci) { - __m128i sraw = _mm_loadl_epi64(reinterpret_cast( - scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); - sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); - } - __m256 blk_acc[4]; - for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); - for (int kk = 0; kk < 16; ++kk) { - const int kb = kblk + kk; - const uint8_t* wbase = blk + (size_t)kb * N2 + c0; - constexpr int PFD = 8; - if (kb + PFD < Kpairs) - _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), - _MM_HINT_T0); - const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); - const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); - for (int ci = 0; ci < nchunk; ++ci) { - __m256i wi = _mm256_cvtepu8_epi32( - _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); - __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); - __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); - blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); - blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); - } - } - for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); - } - for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); - } - for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 - float o = 0.0f; - for (int kb = 0; kb < Kpairs; ++kb) { - const uint8_t byte = blk[(size_t)kb * N2 + c]; - uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; - float sc; - std::memcpy(&sc, &bits, 4); - o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + - e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; - } - out[c] = o; - } -} -#endif - -mxgemv_fn select_mxgemv() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (t >= ISA_AVX512) return mxfp4_gemv_avx512; - if (t >= ISA_AVX2) return mxfp4_gemv_avx2; -#endif - (void)t; - return mxfp4_gemv_scalar; -} - -// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches -// torch.float8_e4m3fn / triton .to(float8e4nv). -inline float e4m3_round(float x) { - const float sign = x < 0.0f ? -1.0f : 1.0f; - const float a = std::fabs(x); - if (a == 0.0f) return 0.0f; - if (a >= 448.0f) return sign * 448.0f; - int e; - std::frexp(a, &e); // a in [2^(e-1), 2^e) - float step = std::ldexp(1.0f, e - 4); - const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) - if (step < min_step) step = min_step; - float r = std::nearbyint(a / step) * step; - if (r > 448.0f) r = 448.0f; - return sign * r; -} - -// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). -inline int ceil_log2_pos(float v) { - uint32_t bits; - std::memcpy(&bits, &v, sizeof(bits)); - const int exp = (int)((bits >> 23) & 0xFF); - const int man = (int)(bits & 0x7FFFFF); - return exp - 127 + (man != 0 ? 1 : 0); -} - -// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). -// bf16->fp32 is exact, so this only reorders -- done once per token/route and -// reused across every output row of the GEMV. -inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { - for (int m = 0; m < K / 2; ++m) { - even[m] = bf16_to_f32(src[2 * m]); - odd[m] = bf16_to_f32(src[2 * m + 1]); - } -} - -// DeepSeek-V4 activation FP8 round-trip (bf16 in/out): per 128-block, -// s = 2^ceil(log2(max(|x|,1e-4)/448)); y = round_e4m3(clamp(x/s,+-448)) * s. -void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K) { - for (int b0 = 0; b0 < K; b0 += 128) { - const int b1 = std::min(K, b0 + 128); - float amax = 1e-4f; - for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); - const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); - const float inv_s = 1.0f / s; - for (int i = b0; i < b1; ++i) { - float q = bf16_to_f32(src[i]) * inv_s; - q = std::min(448.0f, std::max(-448.0f, q)); - dst[i] = f32_to_bf16(e4m3_round(q) * s); - } - } -} - -// --------------------------------- executor --------------------------------- - -struct CpuMoeExecutor; - -struct MoeTask { - CpuMoeExecutor* exec; - int layer_id; - int num_tokens; - const bf16_t* x; // [num_tokens, H] - const int32_t* ids; // [num_tokens, top_k] (raw expert ids; <0 = skip) - const float* w; // [num_tokens, top_k] - bf16_t* y; // [num_tokens, H] -}; - -// Output-row tiling. Small enough to give every worker independent work even at -// batch size 1; large enough to amortize the atomic work-grab. -// -// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read -// every expert weight byte exactly once per token (each output row block is owned -// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the -// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). -// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription -// thrashes the spin-barrier. Deferred (not worth it here / for this workload): -// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would -// only pay off in a grouped/batched (dedup) path. -// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. -// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). -// - NUMA: a single node is assumed. Multi-socket machines would split each -// expert's K dimension per node (banks are already per-row contiguous). -constexpr int IBLK = 32; -constexpr int HBLK = 32; - -// -------------------------------- Q4_0 (W4A8) -------------------------------- -// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed -// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a -// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches -// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload -// path streams. -// -// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 -// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to -// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high -// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD -// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then -// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU -// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. -using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); - -float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - float acc = 0.0f; - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - const uint8_t* q = blk + 2; // 16 nibble bytes - const int8_t* a = aq + (size_t)b * 32; - int isum = 0; - for (int j = 0; j < 16; ++j) { - isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j - isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j - } - acc += fp16_to_f32(dh) * asb[b] * (float)isum; - } - return acc; -} - -#if CPU_MOE_X86 -// fp16 block scale -> fp32 via HW F16C (single value in lane 0). -__attribute__((target("f16c"))) -static inline float q4_scale(uint16_t h) { - return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); -} - -// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low -// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). -__attribute__((target("avx2"))) -static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { - const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); - const __m128i lo = _mm_and_si128(qb, mask); - const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); - return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); -} - -// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + -// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no -// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the -// avx512-tier fallback, since the block dot is 256-bit either way). -__attribute__((target("avx2,fma,f16c"))) -float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - const __m256i ones16 = _mm256_set1_epi16(1); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} - -// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). -__attribute__((target("avx2,avxvnni,fma,f16c"))) -float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { - const __m128i mask = _mm_set1_epi8(0x0F); - const __m256i eight = _mm256_set1_epi8(8); - __m256 accF = _mm256_setzero_ps(); - const int nb = K / 32; - for (int b = 0; b < nb; ++b) { - const uint8_t* blk = w + (size_t)b * 18; - _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); - uint16_t dh; - std::memcpy(&dh, blk, sizeof(dh)); - __m256i wq = q4_unpack32(blk, mask, eight); - __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); - __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) - __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) - __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); - // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and - // reduce at the end (equivalent to scale * block_total). - accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); - } - return hsum256(accF); -} -#endif // CPU_MOE_X86 - -// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to -// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when -// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. -q4dot_fn select_q4dot() { - const IsaTier t = pick_isa(); -#if CPU_MOE_X86 - if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; - if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; -#endif - (void)t; - return q4_0_dot_i8_scalar; -} - -enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4 }; - -// Each ctor pointer arg is the address of a CPU int64 array of length -// num_layers (one base address per layer, built by cpu_executor.py's -// _make_table), not a single flat bank. tbl_at resolves -// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr -// arg 0) resolves to nullptr without dereferencing. -inline const void* tbl_at(const uint64_t* tbl, int layer_id) { - return tbl ? reinterpret_cast(tbl[layer_id]) : nullptr; -} - -struct CpuMoeExecutor { - int num_threads; - int num_layers, num_experts, top_k; - int H, I; - int act, apply_on_input; - int fmt; // WFmt - bool needs_di = false; // pre-deinterleave activations to fp32 (nvfp4/ds_fp4) - // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl - // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up - // table (down_tbl likewise for down); which reinterpretation applies is picked by - // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4). - const uint64_t* gate_up_tbl; // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* down_tbl; // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks - const uint64_t* gu_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales - const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals - const uint64_t* dn_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales - const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals - const uint64_t* gu_bias_tbl; // mxfp4: [E,2I] bf16 biases - const uint64_t* dn_bias_tbl; // mxfp4: [E,H] bf16 biases - float swiglu_alpha; - float swiglu_limit; // +inf == no clamp - dot_fn dot; - nvdot_fn nvdot; - nvi8dot_fn nvi8dot = nullptr; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) - bool use_vnni = false; // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8) - bool use_q4a8 = false; // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant - dsdot_fn dsdot; - mxgemv_fn mxgemv; - q4dot_fn q4dot; - // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU - // (same reference grid), so submit() must not repeat it on the host-callback - // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL - // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving - // it to a captured GPU elementwise kernel removes it while keeping the official - // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). - bool input_prequant = false; - // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). - int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; - float e2m1_lut[16]; - float e4m3_lut[256]; - float e8m0_lut[256]; // mxfp4 block scale: 2^(s-127), s clamped to [0,254] - const char* isa; - - std::vector g_scratch; // [max_tokens * top_k * I] intermediate - std::vector xq_scratch; // [max_tokens * H] ds_fp4 fp8-roundtripped input - // ds_fp4 activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot. - std::vector xe_scratch, xo_scratch; // [max_tokens * H/2] (input) - std::vector ge_scratch, go_scratch; // [max_tokens*top_k*I/2] (intermediate) - // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale. - std::vector xi8_scratch, gi8_scratch; // [max_tokens*H], [max_tokens*top_k*I] - std::vector xas_scratch, gas_scratch; // [max_tokens*H/16], [..*top_k*I/16] - std::string isa_str; - - std::vector workers; - std::mutex task_mtx; - std::condition_variable task_cv; - std::mutex sync_mtx; - std::condition_variable sync_cv; - - bool stop = false; - uint64_t cur_gen = 0; - MoeTask* cur_task = nullptr; - std::atomic submitted{0}; - std::atomic completed{0}; - - std::atomic p1_next{0}; - std::atomic p2_next{0}; - std::atomic prt_next{0}; // ds_fp4 intermediate fp8 round-trip phase - int64_t p1_total = 0, p2_total = 0, prt_total = 0; - int n_iblk = 0, n_hblk = 0; - std::atomic done_count{0}; - std::atomic bar_count{0}; - std::atomic bar_sense{0}; - - std::vector owned_tasks; // persistent task descriptors (graph-stable) - std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) - - // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- - // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls - // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU - // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round - // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a - // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cas…5647 tokens truncated…h size) pair -- the Python side // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: // the same pointers are used by the GPU kernels and by this thread). std::thread coord_thread; diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index 99539d1f2..00a1b540d 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -19,6 +19,14 @@ #include #include +#ifndef CUDART_CB +#define CUDART_CB +#endif + +#ifndef __grid_constant__ +#define __grid_constant__ +#endif + // --- API name mapping (CUDA -> HIP) --- // HIP already defines most cuda* names as macros that expand to hip* equivalents // via hip_runtime.h, but a few are missing or differ in signature. Define them @@ -52,6 +60,14 @@ #define cudaHostAlloc hipHostMalloc #endif +#ifndef cudaHostAllocPortable +#define cudaHostAllocPortable hipHostMallocPortable +#endif + +#ifndef cudaHostAllocMapped +#define cudaHostAllocMapped hipHostMallocMapped +#endif + #ifndef cudaHostRegister #define cudaHostRegister hipHostRegister #endif @@ -122,6 +138,14 @@ #define cudaStream_t hipStream_t #endif +#ifndef cudaStreamSynchronize +#define cudaStreamSynchronize hipStreamSynchronize +#endif + +#ifndef cudaLaunchHostFunc +#define cudaLaunchHostFunc hipLaunchHostFunc +#endif + #ifndef dim3 // HIP already provides dim3; this is a no-op guard. #endif diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index 8e917832c..f21b9585f 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -1,5 +1,6 @@ #pragma once +#include #include #include @@ -115,6 +116,10 @@ public: } auto with_attr(bool use_pdl) -> LaunchKernel & { +#ifdef __HIP__ + (void)use_pdl; + m_config.numAttrs = 0; +#else if (use_pdl) { m_attr_cache.id = ::cudaLaunchAttributeProgrammaticStreamSerialization; m_attr_cache.val.programmaticStreamSerializationAllowed = 1; @@ -123,6 +128,7 @@ public: } else { m_config.numAttrs = 0; } +#endif return *this; } diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index bb83c23ed..fe3f6be49 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -34,40 +34,64 @@ inline constexpr auto get_mem_package() { } __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { +#ifdef __HIP_PLATFORM_AMD__ + return *src; +#else uint32_t tmp; asm volatile("ld.global.L1::no_allocate.b32 %0,[%1];" : "=r"(tmp) : "l"(src)); return uint1{tmp}; +#endif } __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { +#ifdef __HIP_PLATFORM_AMD__ + return *src; +#else uint32_t tmp0, tmp1; asm volatile("ld.global.L1::no_allocate.v2.b32 {%0,%1},[%2];" : "=r"(tmp0), "=r"(tmp1) : "l"(src)); return uint2{tmp0, tmp1}; +#endif } __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { +#ifdef __HIP_PLATFORM_AMD__ + return *src; +#else uint32_t tmp0, tmp1, tmp2, tmp3; asm volatile("ld.global.L1::no_allocate.v4.b32 {%0,%1,%2,%3},[%4];" : "=r"(tmp0), "=r"(tmp1), "=r"(tmp2), "=r"(tmp3) : "l"(src)); return uint4{tmp0, tmp1, tmp2, tmp3}; +#endif } __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { +#ifdef __HIP_PLATFORM_AMD__ + *dst = value; +#else uint32_t tmp = value.x; asm volatile("st.global.wt.b32 [%0],%1;" ::"l"(dst), "r"(tmp)); +#endif } __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { +#ifdef __HIP_PLATFORM_AMD__ + *dst = value; +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; asm volatile("st.global.wt.v2.b32 [%0],{%1,%2};" ::"l"(dst), "r"(tmp0), "r"(tmp1)); +#endif } __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { +#ifdef __HIP_PLATFORM_AMD__ + *dst = value; +#else uint32_t tmp0 = value.x; uint32_t tmp1 = value.y; uint32_t tmp2 = value.z; uint32_t tmp3 = value.w; asm volatile("st.global.wt.v4.b32 [%0],{%1,%2,%3,%4};" ::"l"(dst), "r"(tmp0), "r"(tmp1), "r"(tmp2), "r"(tmp3)); +#endif } __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag_ptr) { diff --git a/python/freetoken/kernel/fast_index_copy.py b/python/freetoken/kernel/fast_index_copy.py index 05b6a4f83..1aaa1303d 100644 --- a/python/freetoken/kernel/fast_index_copy.py +++ b/python/freetoken/kernel/fast_index_copy.py @@ -22,42 +22,6 @@ def _skip_fast_index_copy_enabled() -> bool: return os.getenv(SKIP_FAST_INDEX_COPY_ENV, "").strip().lower() in _TRUE_VALUES -def _is_rocm() -> bool: - return getattr(torch.version, "hip", None) is not None - - -def _rocm_index_copy_fallback( - dst: torch.Tensor, - dst_indices: torch.Tensor, - src: torch.Tensor, - src_indices: torch.Tensor, - num_indices: torch.Tensor | None = None, -) -> None: - """Correctness-first ROCm fallback for the CUDA-specific copy JIT. - - The native fast-index-copy header still contains NVIDIA inline PTX and CUDA - DLPack device matchers. Until that kernel has a HIP implementation, keep - ROCm functional by gathering only the requested source rows and moving that - bounded selection to the destination device. CUDA continues to use the - existing JIT unchanged. - """ - count = dst_indices.numel() if num_indices is None else int(num_indices.item()) - assert 0 <= count <= dst_indices.numel() - assert count <= src_indices.numel() - if count == 0: - return - - src_index = src_indices[:count].to(device=src.device, dtype=torch.long) - dst_index = dst_indices[:count].to(device=dst.device, dtype=torch.long) - rows = src.index_select(0, src_index) - if rows.device != dst.device: - rows = rows.to( - device=dst.device, - non_blocking=src.device.type == "cpu" and src.is_pinned(), - ) - dst.index_copy_(0, dst_index, rows) - - @lru_cache(maxsize=None) def _jit_update_flag_module() -> Module: return load_jit( @@ -150,14 +114,6 @@ def fast_index_copy_jit( dst = dst.as_strided(size=(dst.size(0), num_dst_feature), stride=(num_dst_feature, 1)) src = src.as_strided(size=(src.size(0), num_src_feature), stride=(num_src_feature, 1)) - if _is_rocm(): - if priority is not None: - raise NotImplementedError( - "ROCm fast-index-copy fallback does not implement high/normal priority scheduling" - ) - _rocm_index_copy_fallback(dst, dst_indices, src, src_indices, num_indices) - return - feature_size = dst.size(-1) * dst.element_size() num_block = num_block or DEFAULT_NUM_BLOCKS worker_threads = worker_threads or _default_worker_threads(feature_size) diff --git a/python/freetoken/kernel/triton/activation.py b/python/freetoken/kernel/triton/activation.py index 2c38b533e..0b7c945c3 100644 --- a/python/freetoken/kernel/triton/activation.py +++ b/python/freetoken/kernel/triton/activation.py @@ -20,8 +20,9 @@ import triton.language as tl from triton.language.extra import libdevice from triton.language.extra.cuda import gdc_wait, gdc_launch_dependents +from triton.language import target_info -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported SILU = 0 GELU = 1 @@ -48,6 +49,8 @@ def _pdl_supported() -> bool: @triton.jit def _fast_tanh(x): + if target_info.is_hip(): + return libdevice.tanh(x) # PTX tanh.approx.f32 — single HW op, matches flashinfer math::tanh. return tl.inline_asm_elementwise( "tanh.approx.f32 $0, $1;", "=f,f", [x], @@ -57,6 +60,8 @@ def _fast_tanh(x): @triton.jit def _fast_ex2(x): + if target_info.is_hip(): + return libdevice.exp2(x) # PTX ex2.approx.f32 — matches __expf fast path used by flashinfer silu. return tl.inline_asm_elementwise( "ex2.approx.f32 $0, $1;", "=f,f", [x], @@ -134,8 +139,9 @@ def _act_and_mul( block_d = min(triton.next_power_of_2(d), 1024 if M >= 4096 else 512) num_stages = 2 if block_d == 1024 else 3 _act_and_mul_kernel[grid]( - o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, launch_pdl=pdl, + o2, x2, d, alpha, limit, ACT=kind, ENABLE_PDL=pdl, BLOCK_D=block_d, num_warps=4, num_stages=num_stages, + **({} if is_rocm() else {"launch_pdl": pdl}), ) return out diff --git a/python/freetoken/kernel/triton/e4m3_compat.py b/python/freetoken/kernel/triton/e4m3_compat.py index 61d3a0e77..e52095cc4 100644 --- a/python/freetoken/kernel/triton/e4m3_compat.py +++ b/python/freetoken/kernel/triton/e4m3_compat.py @@ -59,6 +59,10 @@ def e4m3_native() -> bool: if _native is None: if FORCE_EMU: _native = False + elif torch.version.hip is not None: + # ROCm reports gfx1101 as capability (11, 0), which is not a CUDA + # compute capability and must not select the native fp8e4nv path. + _native = False else: native = {torch.cuda.get_device_capability(i) >= (8, 9) for i in range(torch.cuda.device_count())} diff --git a/python/freetoken/kernel/triton/norm.py b/python/freetoken/kernel/triton/norm.py index 3f95c29f0..9071e1dff 100644 --- a/python/freetoken/kernel/triton/norm.py +++ b/python/freetoken/kernel/triton/norm.py @@ -30,7 +30,7 @@ import triton.language as tl from triton.language.extra.cuda import gdc_launch_dependents, gdc_wait -from freetoken.utils.arch import is_sm90_supported +from freetoken.utils.arch import is_rocm, is_sm90_supported _HEUR = {"BLOCK": lambda a: triton.next_power_of_2(a["H"])} @@ -144,7 +144,8 @@ def _rmsnorm(input, weight, eps, out, gemma: bool): pdl = contig and is_sm90_supported() _rmsnorm_kernel[(A, B)]( out, input, weight, eps, H, sxa, sxb, soa, sob, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) return out @@ -172,7 +173,8 @@ def _fused_add_rmsnorm(input, residual, weight, eps, gemma: bool): pdl = contig and is_sm90_supported() _fused_add_rmsnorm_kernel[(A, B)]( input, residual, weight, eps, H, sxa, sxb, sra, srb, - CONTIG=contig, ENABLE_PDL=pdl, launch_pdl=pdl, GEMMA=gemma, + CONTIG=contig, ENABLE_PDL=pdl, GEMMA=gemma, + **({} if is_rocm() else {"launch_pdl": pdl}), num_warps=_num_warps(A * B), num_stages=1, ) diff --git a/setup.py b/setup.py index 8ba0640a0..698ad780c 100644 --- a/setup.py +++ b/setup.py @@ -9,6 +9,7 @@ ROOT = Path(__file__).parent +KERNEL_INCLUDE = ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include" def _check_toolchain() -> None: @@ -61,6 +62,8 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: runtime_lib = "cudart" extra_compile = ["-O3", "-std=c++17"] +runtime_include_dirs.append(str(KERNEL_INCLUDE)) + _check_toolchain() diff --git a/tests/kernels/test_fast_index_copy_rocm.py b/tests/kernels/test_fast_index_copy_rocm.py deleted file mode 100644 index 37f0d5f64..000000000 --- a/tests/kernels/test_fast_index_copy_rocm.py +++ /dev/null @@ -1,64 +0,0 @@ -from __future__ import annotations - -import pytest -import torch - -import freetoken.kernel.fast_index_copy as fast_copy - - -def test_rocm_fallback_copies_only_requested_rows() -> None: - src = torch.arange(20, dtype=torch.float32).reshape(5, 4).to(torch.bfloat16) - dst = torch.full((4, 4), -1, dtype=torch.bfloat16) - src_indices = torch.tensor([4, 1, 3], dtype=torch.int32) - dst_indices = torch.tensor([2, 0, 3], dtype=torch.int32) - num_indices = torch.tensor([2], dtype=torch.int64) - - fast_copy._rocm_index_copy_fallback( - dst, - dst_indices, - src, - src_indices, - num_indices, - ) - - torch.testing.assert_close(dst[2], src[4], rtol=0, atol=0) - torch.testing.assert_close(dst[0], src[1], rtol=0, atol=0) - torch.testing.assert_close(dst[1], torch.full((4,), -1, dtype=torch.bfloat16), rtol=0, atol=0) - torch.testing.assert_close(dst[3], torch.full((4,), -1, dtype=torch.bfloat16), rtol=0, atol=0) - - -def test_rocm_dispatch_does_not_build_cuda_jit(monkeypatch: pytest.MonkeyPatch) -> None: - src = torch.arange(12, dtype=torch.float32).reshape(3, 4) - dst = torch.zeros((3, 4), dtype=torch.float32) - src_indices = torch.tensor([2, 0], dtype=torch.int32) - dst_indices = torch.tensor([1, 2], dtype=torch.int32) - - monkeypatch.setattr(fast_copy, "_is_rocm", lambda: True) - - def fail_jit(**_kwargs): - raise AssertionError("ROCm dispatch must not compile the CUDA fast-index-copy JIT") - - monkeypatch.setattr(fast_copy, "_jit_fast_index_copy_module", fail_jit) - - fast_copy.fast_index_copy_jit(dst, dst_indices, src, src_indices) - - torch.testing.assert_close(dst[1], src[2], rtol=0, atol=0) - torch.testing.assert_close(dst[2], src[0], rtol=0, atol=0) - - -def test_rocm_priority_mode_fails_closed(monkeypatch: pytest.MonkeyPatch) -> None: - src = torch.zeros((2, 4), dtype=torch.float32) - dst = torch.zeros((2, 4), dtype=torch.float32) - indices = torch.tensor([0], dtype=torch.int32) - - monkeypatch.setattr(fast_copy, "_is_rocm", lambda: True) - - with pytest.raises(NotImplementedError, match="priority scheduling"): - fast_copy.fast_index_copy_jit( - dst, - indices, - src, - indices, - priority="high", - sync_flag=torch.zeros((1,), dtype=torch.int32), - ) From 62bb9623eb3c802ad53063b906374d4cdb8d1b30 Mon Sep 17 00:00:00 2001 From: Yuu <206304251+nekomario28@users.noreply.github.com> Date: Mon, 24 Aug 2026 03:28:27 +0900 Subject: [PATCH 06/10] fix(rocm): preserve CPU MoE implementation --- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 484 +++++++++++++++++- 1 file changed, 480 insertions(+), 4 deletions(-) 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 ad397fdd7..56ab93df8 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -1,6 +1,3 @@ -Warning: truncated output (original token count: 25647) -Total output lines: 2150 - // CPU-compute MoE executor for the "cpu" offload backend. // // Decode ships activations to the CPU, computes the routed experts here (reading @@ -829,7 +826,486 @@ float dot_dsfp4_avx512(const uint8_t* packed, const uint8_t* scale, const float* // AVX2: a 32-K block is 16 bytes -> two 8-lane halves (8 even + 8 odd each). __attribute__((target("avx2,fma"))) inline __m256 dsfp4_half_avx2(const uint8_t* pk, const float* xeb, const float* xob, __m256 mag8) { - __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cas…5647 tokens truncated…h size) pair -- the Python side + __m256i wi = _mm256_cvtepu8_epi32(_mm_loadl_epi64(reinterpret_cast(pk))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + return _mm256_fmadd_ps(vlo, _mm256_loadu_ps(xeb), _mm256_mul_ps(vhi, _mm256_loadu_ps(xob))); +} + +__attribute__((target("avx2,fma"))) +float dot_dsfp4_avx2(const uint8_t* packed, const uint8_t* scale, const float* xe, + const float* xo, int K, const float* e2m1, const float* e8m0) { + const __m256 mag8 = _mm256_loadu_ps(e2m1); + __m256 acc0 = _mm256_setzero_ps(), acc1 = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* pk = packed + (size_t)b * 16; + const float* xeb = xe + (size_t)b * 16; + const float* xob = xo + (size_t)b * 16; + const __m256 sc = _mm256_set1_ps(e8m0[scale[b]]); + acc0 = _mm256_fmadd_ps(dsfp4_half_avx2(pk, xeb, xob, mag8), sc, acc0); + acc1 = _mm256_fmadd_ps(dsfp4_half_avx2(pk + 8, xeb + 8, xob + 8, mag8), sc, acc1); + } + return hsum256(_mm256_add_ps(acc0, acc1)); +} +#endif + +dsdot_fn select_dsdot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return dot_dsfp4_avx512; + if (t >= ISA_AVX2) return dot_dsfp4_avx2; +#endif + (void)t; + return dot_dsfp4_scalar; +} + +// ------------------------- mxfp4 (gpt-oss) GEMV ----------------------------- +// Transposed split-K layout: blk[Kpairs, N2] (N innermost), scl[Kpairs/16, N2] +// e8m0 per 32-K. Computes out[c] = sum_kb (E2M1[lo]*x[2kb] + E2M1[hi]*x[2kb+1]) +// * 2^(e8m0-127) for a contiguous column tile (blk/scl already offset to col 0 of +// the tile). Vectorized over N (16 columns / __m512), K stays the outer (cache- +// sequential) loop. Used by both gate_up (K=H) and down (K=I). +using mxgemv_fn = void (*)(float*, const uint8_t*, const uint8_t*, const bf16_t*, int, int, + int, const float*, const float*); + +void mxfp4_gemv_scalar(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + for (int c = 0; c < ncol; ++c) out[c] = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t* w = blk + (size_t)kb * N2; + const uint8_t* s = scl + (size_t)(kb >> 4) * N2; + const float xl = bf16_to_f32(x[2 * kb]); + const float xh = bf16_to_f32(x[2 * kb + 1]); + for (int c = 0; c < ncol; ++c) { + const uint8_t byte = w[c]; + out[c] += (e2m1[byte & 0xF] * xl + e2m1[byte >> 4] * xh) * e8m0[s[c]]; + } + } +} + +#if CPU_MOE_X86 +__attribute__((target("avx512f"))) +void mxfp4_gemv_avx512(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[c]=2^(c-127) computed via bit construction (no gather) + const __m512 lut = _mm512_loadu_ps(e2m1); + const __m512i loma = _mm512_set1_epi32(0xF); + // K-outer / N-inner: each kb cache line is read once and all live column chunks + // (up to 4 -> 64 cols) accumulate from registers, so DRAM/L2 stream the tile once. + int c0 = 0; + for (; c0 + 16 <= ncol; c0 += 64) { + const int nchunk = std::min(4, (ncol - c0) / 16); + __m512 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = 32 K = one scale row + __m512 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadu_si128(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 16)); + sc[ci] = _mm512_castsi512_ps(_mm512_slli_epi32(_mm512_cvtepu8_epi32(sraw), 23)); + } + __m512 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm512_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + // The transposed layout strides K by N2 bytes; prefetch ahead so the strided + // reads are not exposed to DRAM latency (the HW streamer misses big strides). + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m512 xl = _mm512_set1_ps(bf16_to_f32(x[2 * kb])); + const __m512 xh = _mm512_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m512i wi = _mm512_cvtepu8_epi32( + _mm_loadu_si128(reinterpret_cast(wbase + ci * 16))); + __m512 vlo = _mm512_permutexvar_ps(_mm512_and_si512(wi, loma), lut); + __m512 vhi = _mm512_permutexvar_ps(_mm512_and_si512(_mm512_srli_epi32(wi, 4), loma), lut); + blk_acc[ci] = _mm512_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm512_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm512_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm512_storeu_ps(out + c0 + ci * 16, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 16) + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} + +__attribute__((target("avx2,fma"))) +void mxfp4_gemv_avx2(float* out, const uint8_t* blk, const uint8_t* scl, const bf16_t* x, + int Kpairs, int N2, int ncol, const float* e2m1, const float* e8m0) { + (void)e8m0; // e8m0[s]=2^(s-127) built via s<<23 (no gather) + const __m256 mag8 = _mm256_loadu_ps(e2m1); + int c0 = 0; + for (; c0 + 8 <= ncol; c0 += 32) { // up to 4 chunks of 8 = 32 cols + const int nchunk = std::min(4, (ncol - c0) / 8); + __m256 acc[4]; + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_setzero_ps(); + for (int kblk = 0; kblk < Kpairs; kblk += 16) { // 16 K-pairs = one scale row + __m256 sc[4]; + for (int ci = 0; ci < nchunk; ++ci) { + __m128i sraw = _mm_loadl_epi64(reinterpret_cast( + scl + (size_t)(kblk >> 4) * N2 + c0 + ci * 8)); + sc[ci] = _mm256_castsi256_ps(_mm256_slli_epi32(_mm256_cvtepu8_epi32(sraw), 23)); + } + __m256 blk_acc[4]; + for (int ci = 0; ci < nchunk; ++ci) blk_acc[ci] = _mm256_setzero_ps(); + for (int kk = 0; kk < 16; ++kk) { + const int kb = kblk + kk; + const uint8_t* wbase = blk + (size_t)kb * N2 + c0; + constexpr int PFD = 8; + if (kb + PFD < Kpairs) + _mm_prefetch(reinterpret_cast(blk + (size_t)(kb + PFD) * N2 + c0), + _MM_HINT_T0); + const __m256 xl = _mm256_set1_ps(bf16_to_f32(x[2 * kb])); + const __m256 xh = _mm256_set1_ps(bf16_to_f32(x[2 * kb + 1])); + for (int ci = 0; ci < nchunk; ++ci) { + __m256i wi = _mm256_cvtepu8_epi32( + _mm_loadl_epi64(reinterpret_cast(wbase + ci * 8))); + __m256 vlo = e2m1_decode8(_mm256_and_si256(wi, _mm256_set1_epi32(0xF)), mag8); + __m256 vhi = e2m1_decode8(_mm256_srli_epi32(wi, 4), mag8); + blk_acc[ci] = _mm256_fmadd_ps(vlo, xl, blk_acc[ci]); + blk_acc[ci] = _mm256_fmadd_ps(vhi, xh, blk_acc[ci]); + } + } + for (int ci = 0; ci < nchunk; ++ci) acc[ci] = _mm256_fmadd_ps(blk_acc[ci], sc[ci], acc[ci]); + } + for (int ci = 0; ci < nchunk; ++ci) _mm256_storeu_ps(out + c0 + ci * 8, acc[ci]); + } + for (int c = c0; c < ncol; ++c) { // tail columns (< 8); none when ncol%8==0 + float o = 0.0f; + for (int kb = 0; kb < Kpairs; ++kb) { + const uint8_t byte = blk[(size_t)kb * N2 + c]; + uint32_t bits = (uint32_t)scl[(size_t)(kb >> 4) * N2 + c] << 23; + float sc; + std::memcpy(&sc, &bits, 4); + o += (e2m1[byte & 0xF] * bf16_to_f32(x[2 * kb]) + + e2m1[byte >> 4] * bf16_to_f32(x[2 * kb + 1])) * sc; + } + out[c] = o; + } +} +#endif + +mxgemv_fn select_mxgemv() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (t >= ISA_AVX512) return mxfp4_gemv_avx512; + if (t >= ISA_AVX2) return mxfp4_gemv_avx2; +#endif + (void)t; + return mxfp4_gemv_scalar; +} + +// Round a clamped |x|<=448 to nearest float8-e4m3 (RNE), back to fp32. Matches +// torch.float8_e4m3fn / triton .to(float8e4nv). +inline float e4m3_round(float x) { + const float sign = x < 0.0f ? -1.0f : 1.0f; + const float a = std::fabs(x); + if (a == 0.0f) return 0.0f; + if (a >= 448.0f) return sign * 448.0f; + int e; + std::frexp(a, &e); // a in [2^(e-1), 2^e) + float step = std::ldexp(1.0f, e - 4); + const float min_step = std::ldexp(1.0f, -9); // e4m3 subnormal step (2^-9) + if (step < min_step) step = min_step; + float r = std::nearbyint(a / step) * step; + if (r > 448.0f) r = 448.0f; + return sign * r; +} + +// IEEE ceil(log2(v)) for v>0 (matches dsv4 _log2_ceil / fast_round_scale). +inline int ceil_log2_pos(float v) { + uint32_t bits; + std::memcpy(&bits, &v, sizeof(bits)); + const int exp = (int)((bits >> 23) & 0xFF); + const int man = (int)(bits & 0x7FFFFF); + return exp - 127 + (man != 0 ? 1 : 0); +} + +// Split an interleaved bf16 row into fp32 even/odd halves (even[m]=src[2m]). +// bf16->fp32 is exact, so this only reorders -- done once per token/route and +// reused across every output row of the GEMV. +inline void deinterleave_bf16_f32(const bf16_t* src, float* even, float* odd, int K) { + for (int m = 0; m < K / 2; ++m) { + even[m] = bf16_to_f32(src[2 * m]); + odd[m] = bf16_to_f32(src[2 * m + 1]); + } +} + +// DeepSeek-V4 activation FP8 round-trip (bf16 in/out): per 128-block, +// s = 2^ceil(log2(max(|x|,1e-4)/448)); y = round_e4m3(clamp(x/s,+-448)) * s. +void fp8_roundtrip_bf16(const bf16_t* src, bf16_t* dst, int K) { + for (int b0 = 0; b0 < K; b0 += 128) { + const int b1 = std::min(K, b0 + 128); + float amax = 1e-4f; + for (int i = b0; i < b1; ++i) amax = std::max(amax, std::fabs(bf16_to_f32(src[i]))); + const float s = std::ldexp(1.0f, ceil_log2_pos(amax * (1.0f / 448.0f))); + const float inv_s = 1.0f / s; + for (int i = b0; i < b1; ++i) { + float q = bf16_to_f32(src[i]) * inv_s; + q = std::min(448.0f, std::max(-448.0f, q)); + dst[i] = f32_to_bf16(e4m3_round(q) * s); + } + } +} + +// --------------------------------- executor --------------------------------- + +struct CpuMoeExecutor; + +struct MoeTask { + CpuMoeExecutor* exec; + int layer_id; + int num_tokens; + const bf16_t* x; // [num_tokens, H] + const int32_t* ids; // [num_tokens, top_k] (raw expert ids; <0 = skip) + const float* w; // [num_tokens, top_k] + bf16_t* y; // [num_tokens, H] +}; + +// Output-row tiling. Small enough to give every worker independent work even at +// batch size 1; large enough to amortize the atomic work-grab. +// +// Bandwidth notes (Sapphire Rapids 8480+, 13 cores): the two passes already read +// every expert weight byte exactly once per token (each output row block is owned +// by one worker), and x stays hot in L1 across a (token,expert)'s rows -- so the +// kernel is single-read bandwidth-optimal at bs=1 (~205 GB/s vs ~55 GB/s PCIe). +// One worker per *physical* core, pinned, is the sweet spot; SMT oversubscription +// thrashes the spin-barrier. Deferred (not worth it here / for this workload): +// - AMX-bf16: a GEMM tile engine; decode is M=1 GEMV so tiles sit idle. It would +// only pay off in a grouped/batched (dedup) path. +// - expert dedup for bs>1: read each distinct expert once and GEMM its tokens. +// Helps locality+bytes when bs is large; decode batches here are tiny (<=4). +// - NUMA: a single node is assumed. Multi-socket machines would split each +// expert's K dimension per node (banks are already per-row contiguous). +constexpr int IBLK = 32; +constexpr int HBLK = 32; + +// -------------------------------- Q4_0 (W4A8) -------------------------------- +// Native GGUF Q4_0 experts (gemma4 GGUF): per-32 block = fp16 scale d + 16 packed +// bytes; byte j holds element j in its low nibble and j+16 in its high nibble, so a +// block's storage order is [lo0..lo15, hi0..hi15] and w = (nibble - 8) * d. Matches +// the reference dequant (models/gguf/dequant.py) and the packed banks the GPU offload +// path streams. +// +// llama.cpp ggml_vec_dot_q4_0_q8_0: W4A8. The activation is pre-quantized to Q8_0 +// (per-32-block int8 ``aq`` + fp32 scale ``asb``); each block unpacks its 16 bytes to +// 32 int8 weights in [-8,7] (bytes_from_nibbles_32: low nibbles -> elems 0..15, high +// -> 16..31) and runs an integer block dot -- VPDPBUSD (AVX-VNNI) or VPMADDUBSW+VPMADDWD +// (AVX2) with the ggml sign trick |w|*(sign(w)*a)=w*a, or a scalar int loop -- then +// scales the block sum by wd*xd in fp32. No fp weight dequant / shuffle chain. The GPU +// offload path (ggml_moe_a8_vec / MMVQ) is also W4A8, so cpu and hybrid stay close. +using q4dot_fn = float (*)(const uint8_t*, const int8_t*, const float*, int); + +float q4_0_dot_i8_scalar(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + float acc = 0.0f; + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + const uint8_t* q = blk + 2; // 16 nibble bytes + const int8_t* a = aq + (size_t)b * 32; + int isum = 0; + for (int j = 0; j < 16; ++j) { + isum += ((int)(q[j] & 0x0F) - 8) * (int)a[j]; // elem j + isum += ((int)(q[j] >> 4) - 8) * (int)a[16 + j]; // elem 16+j + } + acc += fp16_to_f32(dh) * asb[b] * (float)isum; + } + return acc; +} + +#if CPU_MOE_X86 +// fp16 block scale -> fp32 via HW F16C (single value in lane 0). +__attribute__((target("f16c"))) +static inline float q4_scale(uint16_t h) { + return _mm_cvtss_f32(_mm_cvtph_ps(_mm_cvtsi32_si128((int)h))); +} + +// Unpack one Q4_0 block's 16 bytes -> 32 int8 weights in [-8,7] (elems 0..15 = low +// nibbles, 16..31 = high nibbles). ``eight`` = _mm256_set1_epi8(8). +__attribute__((target("avx2"))) +static inline __m256i q4_unpack32(const uint8_t* blk, __m128i mask, __m256i eight) { + const __m128i qb = _mm_loadu_si128(reinterpret_cast(blk + 2)); + const __m128i lo = _mm_and_si128(qb, mask); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(qb, 4), mask); + return _mm256_sub_epi8(_mm256_set_m128i(hi, lo), eight); +} + +// AVX2 W4A8 (llama.cpp non-VNNI mul_sum_i8_pairs): integer block dot via VPMADDUBSW + +// VPMADDWD (sign trick), scaled by wd*xd. |aw*sa| pair sums <= 8*127*2 < 32767 -> no +// int16 saturation. This is the fast path on AVX2 CPUs without AVX-VNNI (and the +// avx512-tier fallback, since the block dot is 256-bit either way). +__attribute__((target("avx2,fma,f16c"))) +float q4_0_dot_i8_avx2(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + const __m256i ones16 = _mm256_set1_epi16(1); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i d32 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, sa), ones16); // 8 int32 + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d32), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} + +// AVX-VNNI W4A8: one VPDPBUSD per block (the fast path on modern CPUs). +__attribute__((target("avx2,avxvnni,fma,f16c"))) +float q4_0_dot_i8_vnni(const uint8_t* w, const int8_t* aq, const float* asb, int K) { + const __m128i mask = _mm_set1_epi8(0x0F); + const __m256i eight = _mm256_set1_epi8(8); + __m256 accF = _mm256_setzero_ps(); + const int nb = K / 32; + for (int b = 0; b < nb; ++b) { + const uint8_t* blk = w + (size_t)b * 18; + _mm_prefetch(reinterpret_cast(blk) + 512, _MM_HINT_T0); + uint16_t dh; + std::memcpy(&dh, blk, sizeof(dh)); + __m256i wq = q4_unpack32(blk, mask, eight); + __m256i a = _mm256_loadu_si256(reinterpret_cast(aq + (size_t)b * 32)); + __m256i aw = _mm256_sign_epi8(wq, wq); // |wq| (unsigned operand) + __m256i sa = _mm256_sign_epi8(a, wq); // sign(wq) * a (signed operand) + __m256i di = _mm256_dpbusd_avx_epi32(_mm256_setzero_si256(), aw, sa); + // All 32 elems of the block share wd*xd; distribute over di's 8 partial sums and + // reduce at the end (equivalent to scale * block_total). + accF = _mm256_fmadd_ps(_mm256_cvtepi32_ps(di), _mm256_set1_ps(q4_scale(dh) * asb[b]), accF); + } + return hsum256(accF); +} +#endif // CPU_MOE_X86 + +// All tiers are W4A8 (int8 activations pre-quantized to Q8_0). AVX-VNNI is orthogonal to +// the ISA tier (gated by cpu_has_avxvnni() / FREETOKEN_CPU_MOE_NO_VNNI), so it wins when +// present; otherwise the 256-bit VPMADDUBSW kernel covers both the avx2 and avx512 tiers. +q4dot_fn select_q4dot() { + const IsaTier t = pick_isa(); +#if CPU_MOE_X86 + if (cpu_has_avxvnni()) return q4_0_dot_i8_vnni; + if (t >= ISA_AVX2) return q4_0_dot_i8_avx2; +#endif + (void)t; + return q4_0_dot_i8_scalar; +} + +enum WFmt { WF_BF16 = 0, WF_NVFP4 = 1, WF_MXFP4 = 2, WF_DSFP4 = 3, WF_Q4_0 = 4 }; + +// Each ctor pointer arg is the address of a CPU int64 array of length +// num_layers (one base address per layer, built by cpu_executor.py's +// _make_table), not a single flat bank. tbl_at resolves +// tbl[layer_id] once per task/pass; a null table (bank unused by this fmt, ptr +// arg 0) resolves to nullptr without dereferencing. +inline const void* tbl_at(const uint64_t* tbl, int layer_id) { + return tbl ? reinterpret_cast(tbl[layer_id]) : nullptr; +} + +struct CpuMoeExecutor { + int num_threads; + int num_layers, num_experts, top_k; + int H, I; + int act, apply_on_input; + int fmt; // WFmt + bool needs_di = false; // pre-deinterleave activations to fp32 (nvfp4/ds_fp4) + // Per-layer pointer tables (one base address per layer, see tbl_at). gate_up_tbl + // doubles as the bf16 gate_up table and the nvfp4/mxfp4/q4_0/ds_fp4 packed-gate_up + // table (down_tbl likewise for down); which reinterpretation applies is picked by + // fmt at each resolve site (see gemm1_dot/gemm2_dot/do_pass1_mxfp4/do_pass1_dsfp4). + const uint64_t* gate_up_tbl; // bf16: [E,2I,H] rows; else: packed e2m1/mxfp4-blocks + const uint64_t* down_tbl; // bf16: [E,H,I] rows; else: packed e2m1/mxfp4-blocks + const uint64_t* gu_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,2I,*] block scales + const uint64_t* gu_global_tbl; // nvfp4: [E,2I] fp16 row globals + const uint64_t* dn_scale_tbl; // nvfp4/mxfp4/ds_fp4: [E,H,*] block scales + const uint64_t* dn_global_tbl; // nvfp4: [E,H] fp16 row globals + const uint64_t* gu_bias_tbl; // mxfp4: [E,2I] bf16 biases + const uint64_t* dn_bias_tbl; // mxfp4: [E,H] bf16 biases + float swiglu_alpha; + float swiglu_limit; // +inf == no clamp + dot_fn dot; + nvdot_fn nvdot; + nvi8dot_fn nvi8dot = nullptr; // AVX-VNNI W4A8 nvfp4 dot (nullptr -> use fp32 nvdot) + bool use_vnni = false; // nvfp4 + AVX-VNNI: decode via int8 VPDPBUSD (W4A8) + bool use_q4a8 = false; // q4_0: always W4A8 (llama.cpp Q4_0 x Q8_0); int8 pre-quant + dsdot_fn dsdot; + mxgemv_fn mxgemv; + q4dot_fn q4dot; + // ds_fp4: the caller already FP8-round-tripped the input activations on the GPU + // (same reference grid), so submit() must not repeat it on the host-callback + // thread. That scalar per-element pass is single-threaded ON THE DECODE CRITICAL + // PATH (~0.3ms/layer at H=4096, every worker and the GPU waiting on it); moving + // it to a captured GPU elementwise kernel removes it while keeping the official + // W4A8 numerics bit-exact. Set via set_input_prequant (see cpu_executor.py). + bool input_prequant = false; + // Q4_0 packed-row byte strides (H/32*18 for gate_up over K=H, I/32*18 for down over K=I). + int q4_gu_row_bytes = 0, q4_dn_row_bytes = 0; + float e2m1_lut[16]; + float e4m3_lut[256]; + float e8m0_lut[256]; // mxfp4 block scale: 2^(s-127), s clamped to [0,254] + const char* isa; + + std::vector g_scratch; // [max_tokens * top_k * I] intermediate + std::vector xq_scratch; // [max_tokens * H] ds_fp4 fp8-roundtripped input + // ds_fp4 activations pre-deinterleaved to fp32 (even/odd K) for the row-major dot. + std::vector xe_scratch, xo_scratch; // [max_tokens * H/2] (input) + std::vector ge_scratch, go_scratch; // [max_tokens*top_k*I/2] (intermediate) + // AVX-VNNI W4A8: per-16-block int8 activations [even(8),odd(8)] + per-block scale. + std::vector xi8_scratch, gi8_scratch; // [max_tokens*H], [max_tokens*top_k*I] + std::vector xas_scratch, gas_scratch; // [max_tokens*H/16], [..*top_k*I/16] + std::string isa_str; + + std::vector workers; + std::mutex task_mtx; + std::condition_variable task_cv; + std::mutex sync_mtx; + std::condition_variable sync_cv; + + bool stop = false; + uint64_t cur_gen = 0; + MoeTask* cur_task = nullptr; + std::atomic submitted{0}; + std::atomic completed{0}; + + std::atomic p1_next{0}; + std::atomic p2_next{0}; + std::atomic prt_next{0}; // ds_fp4 intermediate fp8 round-trip phase + int64_t p1_total = 0, p2_total = 0, prt_total = 0; + int n_iblk = 0, n_hblk = 0; + std::atomic done_count{0}; + std::atomic bar_count{0}; + std::atomic bar_sense{0}; + + std::vector owned_tasks; // persistent task descriptors (graph-stable) + std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) + + // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- + // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls + // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU + // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round + // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a + // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: // the same pointers are used by the GPU kernels and by this thread). std::thread coord_thread; From c7a95f726b3d67f7570668d7140a258253160de5 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Mon, 24 Aug 2026 15:08:37 +0800 Subject: [PATCH 07/10] build(rocm): support RDNA4 discovery and modular SDKs --- docs/install.md | 31 +++++++- python/freetoken/kernel/__main__.py | 9 ++- .../csrc/include/freetoken/hip_compat.h | 49 +++++++------ python/freetoken/kernel/utils.py | 72 +++++++++++++++++-- python/freetoken/utils/__init__.py | 2 + python/freetoken/utils/arch.py | 46 ++++++++++-- setup.py | 65 +++++++++++------ tests/kernels/test_pinned_tensor.py | 7 +- tests/utils/test_rocm_arch.py | 53 ++++++++++++++ 9 files changed, 271 insertions(+), 63 deletions(-) create mode 100644 tests/utils/test_rocm_arch.py diff --git a/docs/install.md b/docs/install.md index f5205ab3b..4db8a15a3 100644 --- a/docs/install.md +++ b/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 diff --git a/python/freetoken/kernel/__main__.py b/python/freetoken/kernel/__main__.py index 5227443e3..457385a58 100644 --- a/python/freetoken/kernel/__main__.py +++ b/python/freetoken/kernel/__main__.py @@ -6,7 +6,7 @@ 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__) @@ -14,7 +14,9 @@ def generate_clangd(): include_paths = [find_include_path(), find_dlpack_include_path()] + DEFAULT_INCLUDE # TODO(ROCm): hiprtc JIT cache should be separate from nvcc JIT cache to avoid stale binaries. - try: + if is_rocm(): + arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"] + else: status = subprocess.run( args=["nvidia-smi", "--query-gpu=compute_cap", "--format=csv,noheader"], capture_output=True, @@ -23,9 +25,6 @@ def generate_clangd(): compute_cap = status.stdout.decode("utf-8").strip().split("\n")[0] major, minor = compute_cap.split(".") arch_flags = ["-xcuda", f"--cuda-gpu-arch=sm_{major}{minor}"] - except (subprocess.CalledProcessError, FileNotFoundError): - # TODO(ROCm): parse gfx target from rocm-smi; default to gfx1100 for now. - arch_flags = ["-xhip", "--offload-arch=gfx1100"] compile_flags = ",\n ".join( arch_flags + ["-std=c++20", "-Wall", "-Wextra"] + [f"-isystem{path}" for path in include_paths] diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index 00a1b540d..eb7e71e18 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -7,26 +7,22 @@ // On NVIDIA platforms the CUDA headers are included as-is and every macro below // resolves to the original CUDA symbol, so there is zero overhead. // -// Supported ROCm targets (RDNA3): +// Supported ROCm targets: // gfx1100 — RX 7900 XTX / XT // gfx1101 — RX 7900 GRE // gfx1102 — RX 7700 / XT // gfx1103 — RX 7600 / XT +// gfx1200 — RX 9060 family +// gfx1201 — RX 9070 family / Radeon AI PRO R9700 -#ifdef __HIP__ +#if defined(__HIP_PLATFORM_AMD__) || defined(USE_ROCM) + +#define FREETOKEN_USE_ROCM 1 // --- HIP runtime headers --- #include #include -#ifndef CUDART_CB -#define CUDART_CB -#endif - -#ifndef __grid_constant__ -#define __grid_constant__ -#endif - // --- API name mapping (CUDA -> HIP) --- // HIP already defines most cuda* names as macros that expand to hip* equivalents // via hip_runtime.h, but a few are missing or differ in signature. Define them @@ -60,14 +56,6 @@ #define cudaHostAlloc hipHostMalloc #endif -#ifndef cudaHostAllocPortable -#define cudaHostAllocPortable hipHostMallocPortable -#endif - -#ifndef cudaHostAllocMapped -#define cudaHostAllocMapped hipHostMallocMapped -#endif - #ifndef cudaHostRegister #define cudaHostRegister hipHostRegister #endif @@ -80,6 +68,14 @@ #define cudaHostRegisterMapped hipHostRegisterMapped #endif +#ifndef cudaHostAllocPortable +#define cudaHostAllocPortable hipHostMallocPortable +#endif + +#ifndef cudaHostAllocMapped +#define cudaHostAllocMapped hipHostMallocMapped +#endif + #ifndef cudaHostGetDevicePointer #define cudaHostGetDevicePointer hipHostGetDevicePointer #endif @@ -115,8 +111,7 @@ #endif #ifndef cudaLaunchKernelEx -// TODO(ROCm): hipLaunchKernelEx is available in newer ROCm; use it when widely shipped. -// For now, fall back to hipLaunchKernel with config extracted manually. +// ROCm 7 exposes the CUDA-compatible extended launch configuration through HIP. #define cudaLaunchKernelEx hipLaunchKernelEx #endif @@ -146,12 +141,22 @@ #define cudaLaunchHostFunc hipLaunchHostFunc #endif +#ifndef CUDART_CB +#define CUDART_CB +#endif + +#ifndef __grid_constant__ +#define __grid_constant__ +#endif + #ifndef dim3 // HIP already provides dim3; this is a no-op guard. #endif -#else // !__HIP__ — NVIDIA CUDA path +#else // NVIDIA CUDA path + +#define FREETOKEN_USE_ROCM 0 #include -#endif // __HIP__ +#endif diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index dfdcd4c7b..937a750cb 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -4,6 +4,7 @@ import os import pathlib import re +from functools import cache from typing import TYPE_CHECKING, List, NamedTuple, Tuple, TypeAlias, Union if TYPE_CHECKING: @@ -21,6 +22,7 @@ DEFAULT_CUDA_CFLAGS = ["-std=c++20", "-O3", "--expt-relaxed-constexpr"] DEFAULT_HIP_CFLAGS = ["-std=c++20", "-O3"] DEFAULT_LDFLAGS = [] +DEFAULT_ROCM_ARCHES = ("gfx1100", "gfx1101", "gfx1102", "gfx1103", "gfx1200", "gfx1201") def _is_rocm() -> bool: @@ -50,11 +52,65 @@ def _rank(a: str) -> int: def _hip_cflags(extra: List[str]) -> List[str]: """HIP flags for a kernel build on ROCm.""" - # TODO(ROCm): Triton autotune configs need RDNA3-specific tuning (wave count, LDS size). + # TODO(ROCm): Triton autotune configs need RDNA-specific tuning (wave count, LDS size). flags = DEFAULT_HIP_CFLAGS + extra - rocm_arch = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103") - flags = flags + [f"--offload-arch={rocm_arch}"] - return flags + raw_arches = os.getenv("FREETOKEN_ROCM_ARCH") or os.getenv("PYTORCH_ROCM_ARCH", "") + arches = list(dict.fromkeys(re.findall(r"gfx\d+[a-z]?", raw_arches.lower()))) + if not arches: + from freetoken.utils.arch import get_rocm_gfx_arch + + detected = get_rocm_gfx_arch() + arches = [detected] if detected else list(DEFAULT_ROCM_ARCHES) + return flags + [f"--offload-arch={arch}" for arch in arches] + + +@cache +def _rocm_link_flags() -> List[str]: + """Make ROCm's runtime library discoverable to JIT link commands. + + Traditional ROCm installs provide ``libamdhip64.so`` under ``$ROCM_HOME/lib``. + ROCm 7.14 Python SDK images only provide the versioned soname, while TVM-FFI + still links with ``-lamdhip64``. Supply a cache-local unversioned symlink via + ``LIBRARY_PATH`` without modifying the image's Python environment. + """ + candidates: list[pathlib.Path] = [] + if os.getenv("ROCM_HOME"): + candidates.append(pathlib.Path(os.environ["ROCM_HOME"])) + try: + from torch.utils.cpp_extension import ROCM_HOME + + if ROCM_HOME: + candidates.append(pathlib.Path(ROCM_HOME)) + except ImportError: + pass + spec = importlib.util.find_spec("_rocm_sdk_core") + if spec and spec.submodule_search_locations: + candidates.append(pathlib.Path(next(iter(spec.submodule_search_locations)))) + candidates.append(pathlib.Path("/opt/rocm")) + + for rocm_home in dict.fromkeys(candidates): + library_dir = rocm_home / "lib" + unversioned = library_dir / "libamdhip64.so" + link_dir = library_dir + if not unversioned.exists(): + versioned = sorted(library_dir.glob("libamdhip64.so.*")) + if not versioned: + continue + link_dir = pathlib.Path.home() / ".cache" / "freetoken" / "rocm-lib" + link_dir.mkdir(parents=True, exist_ok=True) + compat_link = link_dir / "libamdhip64.so" + if not compat_link.exists() and not compat_link.is_symlink(): + compat_link.symlink_to(versioned[-1]) + + current = [path for path in os.getenv("LIBRARY_PATH", "").split(":") if path] + os.environ["LIBRARY_PATH"] = ":".join( + dict.fromkeys([str(link_dir), str(library_dir), *current]) + ) + return [f"-Wl,-rpath,{library_dir}"] + + raise RuntimeError("Unable to locate libamdhip64 for ROCm JIT linking") + + CPP_TEMPLATE_TYPE: TypeAlias = Union[int, float, bool] @@ -234,8 +290,10 @@ def load_aot( if _is_rocm(): cuda_cflags = _hip_cflags(extra_cuda_cflags) + runtime_ldflags = _rocm_link_flags() else: cuda_cflags = _cuda_cflags(extra_cuda_cflags) + runtime_ldflags = [] return load( name, @@ -243,7 +301,7 @@ def load_aot( cuda_files=cuda_files, extra_cflags=DEFAULT_CFLAGS + extra_cflags, extra_cuda_cflags=cuda_cflags, - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) @@ -294,8 +352,10 @@ def load_jit( if _is_rocm(): cuda_cflags = _hip_cflags(extra_cuda_cflags) + runtime_ldflags = _rocm_link_flags() else: cuda_cflags = _cuda_cflags(extra_cuda_cflags) + runtime_ldflags = [] return load_inline( name, @@ -303,7 +363,7 @@ def load_jit( cuda_sources=cuda_sources, extra_cflags=DEFAULT_CFLAGS + extra_cflags, extra_cuda_cflags=cuda_cflags, - extra_ldflags=DEFAULT_LDFLAGS + extra_ldflags, + extra_ldflags=DEFAULT_LDFLAGS + runtime_ldflags + extra_ldflags, extra_include_paths=DEFAULT_INCLUDE + extra_include_paths, build_directory=build_directory, ) diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index 1348a5286..c54579a16 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -3,6 +3,7 @@ is_rocm, get_rocm_gfx_arch, is_gfx11xx_family, + is_gfx12xx_family, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -41,6 +42,7 @@ "is_rocm", "get_rocm_gfx_arch", "is_gfx11xx_family", + "is_gfx12xx_family", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index a70c8aa39..3bf8fb61e 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -2,9 +2,18 @@ import functools import os +import re from typing import Tuple +_GFX_ARCH_RE = re.compile(r"gfx\d+[a-z]?") + + +def _gfx_arch_from(value: object) -> str | None: + match = _GFX_ARCH_RE.search(str(value).lower()) + return match.group(0) if match else None + + @functools.cache def is_rocm() -> bool: """True when torch is built for ROCm (AMD GPU) instead of CUDA.""" @@ -14,15 +23,31 @@ def is_rocm() -> bool: @functools.cache def get_rocm_gfx_arch() -> str | None: - """The gfx target of the current AMD GPU (e.g. \"gfx1100\"), or None.""" + """Return the current AMD GPU target (for example ``gfx1201``). + + Prefer the runtime device because build variables may contain multiple + semicolon-separated targets. Environment variables remain useful for + cross-compilation and systems where no GPU is currently visible. + """ if not is_rocm(): return None - # TODO(ROCm): parse rocm-smi for auto-detection; for now rely on env. - for env_var in ("PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"): - val = os.getenv(env_var, "") - for gfx in ("gfx1100", "gfx1101", "gfx1102", "gfx1103"): - if gfx in val: - return gfx + + import torch + + if torch.cuda.is_available(): + try: + props = torch.cuda.get_device_properties(torch.cuda.current_device()) + for attr in ("gcnArchName", "arch"): + arch = _gfx_arch_from(getattr(props, attr, "")) + if arch: + return arch + except (AttributeError, RuntimeError): + pass + + for env_var in ("FREETOKEN_ROCM_ARCH", "PYTORCH_ROCM_ARCH", "HCC_AMDGPU_TARGET"): + arch = _gfx_arch_from(os.getenv(env_var, "")) + if arch: + return arch return None @@ -33,6 +58,13 @@ def is_gfx11xx_family() -> bool: return arch is not None and arch.startswith("gfx110") +@functools.cache +def is_gfx12xx_family() -> bool: + """True when the current AMD GPU is RDNA4 (gfx120x).""" + arch = get_rocm_gfx_arch() + return arch is not None and arch.startswith("gfx120") + + @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: import torch diff --git a/setup.py b/setup.py index 698ad780c..bac07c6fb 100644 --- a/setup.py +++ b/setup.py @@ -5,11 +5,11 @@ 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 = ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include" +KERNEL_INCLUDE = str(ROOT / "python" / "freetoken" / "kernel" / "csrc" / "include") def _check_toolchain() -> None: @@ -25,15 +25,36 @@ def _is_rocm() -> bool: return getattr(torch.version, "hip", None) is not None -def _rocm_paths() -> tuple[list[str], list[str]]: - rocm_home = Path(os.getenv("ROCM_HOME", "/opt/rocm")) - if not rocm_home.exists(): - raise RuntimeError( - "ROCM_HOME is required to build on ROCm. Set ROCM_HOME to your ROCm install." - ) - include_dirs = [str(rocm_home / "include")] - library_dirs = [str(rocm_home / "lib")] - return include_dirs, library_dirs +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)) + + # ROCm 7.14 PyTorch images ship the SDK as a Python package instead of + # installing it at /opt/rocm. + 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 (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]]: @@ -52,18 +73,18 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: IS_ROCM = _is_rocm() if IS_ROCM: - runtime_include_dirs, runtime_library_dirs = _rocm_paths() - runtime_lib = "amdhip64" - # TODO(ROCm): allow override via FREETOKEN_ROCM_ARCH; default to all RDNA3. - rocm_arch = os.getenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1101;gfx1102;gfx1103") - extra_compile = ["-O3", "-std=c++17", f"--offload-arch={rocm_arch}"] + runtime_include_dirs, runtime_library_dirs, runtime_lib = _rocm_paths() + runtime_link_args = [f"-Wl,-rpath,{runtime_library_dirs[0]}"] + # These extensions contain host code only. BuildExtension supplies the ROCm + # platform defines to the C++ compiler; offload architecture flags belong on + # HIP device sources and would be rejected by the host compiler here. + extra_compile = ["-O3", "-std=c++17"] else: runtime_include_dirs, runtime_library_dirs = _cuda_runtime_paths() runtime_lib = "cudart" + runtime_link_args = [] extra_compile = ["-O3", "-std=c++17"] -runtime_include_dirs.append(str(KERNEL_INCLUDE)) - _check_toolchain() @@ -74,12 +95,13 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/pinned_tensor.cpp", ], - include_dirs=runtime_include_dirs, + include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], library_dirs=runtime_library_dirs, libraries=[runtime_lib], extra_compile_args=extra_compile, + extra_link_args=runtime_link_args, ), - # CPU-compute MoE executor for --moe-backend cpu. Links cudart/hip for the + # CPU-compute MoE executor for --moe-backend cpu. Links cudart/amdhip64 for the # cudaLaunchHostFunc submit/sync graph nodes; the bf16 GEMV microkernels # use per-function target attributes (avx512bf16/avx512f) + a runtime # __builtin_cpu_supports dispatch, so the single binary stays portable @@ -89,10 +111,11 @@ def _cuda_runtime_paths() -> tuple[list[str], list[str]]: sources=[ "python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp", ], - include_dirs=runtime_include_dirs, + include_dirs=[KERNEL_INCLUDE, *runtime_include_dirs], library_dirs=runtime_library_dirs, libraries=[runtime_lib], extra_compile_args=extra_compile + ["-pthread"], + extra_link_args=runtime_link_args, ), ], cmdclass={"build_ext": BuildExtension.with_options(use_ninja=True)}, diff --git a/tests/kernels/test_pinned_tensor.py b/tests/kernels/test_pinned_tensor.py index e61108fd5..2fee4f25c 100644 --- a/tests/kernels/test_pinned_tensor.py +++ b/tests/kernels/test_pinned_tensor.py @@ -122,7 +122,12 @@ def test_host_device_ptr_is_identity_under_uva(): pytest.skip("non-UVA platform: host_device_ptr rejects unregistered memory instead") # Under UVA cudaHostGetDevicePointer degenerates to identity for any host pointer # (no registration validation); rejection of pageable memory only exists on - # non-identity platforms (Windows/WDDM), where the translation is real. + # non-identity CUDA platforms (Windows/WDDM), where the translation is real. + # HIP validates registration even though registered/pinned memory uses the + # identity address on Linux. Calling it with pageable memory also leaves a + # sticky HIP error, so the pinned identity case above is the relevant check. + if torch.version.hip is not None: + return pageable = torch.empty(64, dtype=torch.uint8) ext = _load_pinned_extension() assert ext.host_device_ptr(pageable.data_ptr()) == pageable.data_ptr() diff --git a/tests/utils/test_rocm_arch.py b/tests/utils/test_rocm_arch.py new file mode 100644 index 000000000..f080b1e93 --- /dev/null +++ b/tests/utils/test_rocm_arch.py @@ -0,0 +1,53 @@ +from types import SimpleNamespace + +import torch + +from freetoken.utils import arch + + +def _clear_arch_caches() -> None: + arch.get_rocm_gfx_arch.cache_clear() + arch.is_gfx11xx_family.cache_clear() + arch.is_gfx12xx_family.cache_clear() + + +def test_rocm_arch_prefers_visible_device_over_multi_arch_build_env(monkeypatch): + monkeypatch.setattr(arch, "is_rocm", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "current_device", lambda: 0) + monkeypatch.setattr( + torch.cuda, + "get_device_properties", + lambda _device: SimpleNamespace(gcnArchName="gfx1201:sramecc-:xnack-"), + ) + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1100;gfx1200") + _clear_arch_caches() + + assert arch.get_rocm_gfx_arch() == "gfx1201" + assert arch.is_gfx12xx_family() + assert not arch.is_gfx11xx_family() + + _clear_arch_caches() + + +def test_rocm_arch_falls_back_to_cross_compile_env(monkeypatch): + monkeypatch.setattr(arch, "is_rocm", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201") + _clear_arch_caches() + + assert arch.get_rocm_gfx_arch() == "gfx1200" + + _clear_arch_caches() + + +def test_hip_cflags_emit_one_offload_flag_per_arch(monkeypatch): + from freetoken.kernel.utils import _hip_cflags + + monkeypatch.setenv("FREETOKEN_ROCM_ARCH", "gfx1200;gfx1201") + + flags = _hip_cflags(["-Wno-unused-command-line-argument"]) + + assert "--offload-arch=gfx1200" in flags + assert "--offload-arch=gfx1201" in flags + assert not any(";" in flag for flag in flags) From 45417560beeba896a9b0866e50154d2a7a320b8b Mon Sep 17 00:00:00 2001 From: zihaomu Date: Mon, 24 Aug 2026 15:21:29 +0800 Subject: [PATCH 08/10] fix(rocm): complete modular SDK runtime support --- pyproject.toml | 5 ++- python/freetoken/kernel/__main__.py | 19 ++++++---- .../kernel/csrc/jit/fast_index_copy.cuh | 32 ++++++++--------- python/freetoken/kernel/utils.py | 14 ++++---- tests/utils/test_rocm_arch.py | 36 +++++++++++++++++++ 5 files changed, 75 insertions(+), 31 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 3c0c45ede..d22ae67cc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -58,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", ] diff --git a/python/freetoken/kernel/__main__.py b/python/freetoken/kernel/__main__.py index 457385a58..5b66d484d 100644 --- a/python/freetoken/kernel/__main__.py +++ b/python/freetoken/kernel/__main__.py @@ -17,13 +17,18 @@ def generate_clangd(): if is_rocm(): arch_flags = ["-xhip", f"--offload-arch={get_rocm_gfx_arch() or 'gfx1201'}"] else: - 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(".") + 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( arch_flags + ["-std=c++20", "-Wall", "-Wextra"] diff --git a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh index fe3f6be49..bf313c525 100644 --- a/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh +++ b/python/freetoken/kernel/csrc/jit/fast_index_copy.cuh @@ -34,7 +34,7 @@ inline constexpr auto get_mem_package() { } __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 { -#ifdef __HIP_PLATFORM_AMD__ +#if FREETOKEN_USE_ROCM return *src; #else uint32_t tmp; @@ -44,7 +44,7 @@ __always_inline __device__ auto load_nc(const uint1* __restrict__ src) -> uint1 } __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 { -#ifdef __HIP_PLATFORM_AMD__ +#if FREETOKEN_USE_ROCM return *src; #else uint32_t tmp0, tmp1; @@ -54,7 +54,7 @@ __always_inline __device__ auto load_nc(const uint2* __restrict__ src) -> uint2 } __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 { -#ifdef __HIP_PLATFORM_AMD__ +#if FREETOKEN_USE_ROCM return *src; #else uint32_t tmp0, tmp1, tmp2, tmp3; @@ -64,7 +64,7 @@ __always_inline __device__ auto load_nc(const uint4* __restrict__ src) -> uint4 } __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& value) { -#ifdef __HIP_PLATFORM_AMD__ +#if FREETOKEN_USE_ROCM *dst = value; #else uint32_t tmp = value.x; @@ -73,7 +73,7 @@ __always_inline __device__ void store_nc(uint1* __restrict__ dst, const uint1& v } __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& value) { -#ifdef __HIP_PLATFORM_AMD__ +#if FREETOKEN_USE_ROCM *dst = value; #else uint32_t tmp0 = value.x; @@ -83,7 +83,7 @@ __always_inline __device__ void store_nc(uint2* __restrict__ dst, const uint2& v } __always_inline __device__ void store_nc(uint4* __restrict__ dst, const uint4& value) { -#ifdef __HIP_PLATFORM_AMD__ +#if FREETOKEN_USE_ROCM *dst = value; #else uint32_t tmp0 = value.x; @@ -99,7 +99,7 @@ __always_inline __device__ void wait_flag_clear(const int32_t* __restrict__ flag auto* flag = reinterpret_cast(const_cast(flag_ptr)); uint32_t sleep_ns = 128; while (atomicAdd(flag, 0) > 0) { -#if __CUDA_ARCH__ >= 700 +#if !FREETOKEN_USE_ROCM && __CUDA_ARCH__ >= 700 __nanosleep(sleep_ns); #endif sleep_ns = sleep_ns < 2048 ? (sleep_ns << 1) : 2048; @@ -171,7 +171,7 @@ inline bool host_ptr_identity() { } inline void* device_alias(void* ptr, DLDevice dev) { - if (dev.device_type == kDLCUDA || host_ptr_identity()) { + if (dev.device_type == kDLCUDA || dev.device_type == kDLROCM || host_ptr_identity()) { return ptr; } void* mapped = nullptr; @@ -293,7 +293,7 @@ inline auto get_sync_flag_ptr( auto flag_dtype = host::SymbolicDType{}; host::TensorMatcher({1}) .with_dtype(flag_dtype) - .with_device(device) + .with_device(device) .verify(sync_flag); return static_cast(sync_flag.data_ptr()); } @@ -368,17 +368,17 @@ struct FastIndexCopyKernel { TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(src); TensorMatcher({-1, D}) .with_dtype(data_dtype) - .with_device() + .with_device() .verify(dst); TensorMatcher({L}) .with_dtype(indices_dtype) - .with_device(device) + .with_device(device) .verify(src_indices) .verify(dst_indices); @@ -387,7 +387,7 @@ struct FastIndexCopyKernel { const auto num_indices_tensor = num_indices.value(); TensorMatcher({1}) .with_dtype(num_indices_dtype) - .with_device(device) + .with_device(device) .verify(num_indices_tensor); num_indices_data_ptr = static_cast(num_indices_tensor.data_ptr()); @@ -553,14 +553,14 @@ struct MultiIndexCopyKernel { auto indices_dtype = SymbolicDType{}; auto num_indices_dtype = SymbolicDType{}; - TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) + TensorMatcher({B}).with_dtype(ptr_dtype).with_device(device) .verify(dst_ptrs).verify(src_ptrs).verify(feat_bytes); - TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) + TensorMatcher({L}).with_dtype(indices_dtype).with_device(device) .verify(dst_indices).verify(src_indices); const int64_t* valid_length = nullptr; if (num_indices.has_value()) { - TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) + TensorMatcher({1}).with_dtype(num_indices_dtype).with_device(device) .verify(num_indices.value()); valid_length = static_cast(num_indices.value().data_ptr()); } diff --git a/python/freetoken/kernel/utils.py b/python/freetoken/kernel/utils.py index 937a750cb..b59a588f6 100644 --- a/python/freetoken/kernel/utils.py +++ b/python/freetoken/kernel/utils.py @@ -71,7 +71,7 @@ def _rocm_link_flags() -> List[str]: Traditional ROCm installs provide ``libamdhip64.so`` under ``$ROCM_HOME/lib``. ROCm 7.14 Python SDK images only provide the versioned soname, while TVM-FFI still links with ``-lamdhip64``. Supply a cache-local unversioned symlink via - ``LIBRARY_PATH`` without modifying the image's Python environment. + an explicit linker search path without modifying the Python environment. """ candidates: list[pathlib.Path] = [] if os.getenv("ROCM_HOME"): @@ -100,13 +100,13 @@ def _rocm_link_flags() -> List[str]: link_dir.mkdir(parents=True, exist_ok=True) compat_link = link_dir / "libamdhip64.so" if not compat_link.exists() and not compat_link.is_symlink(): - compat_link.symlink_to(versioned[-1]) + try: + compat_link.symlink_to(versioned[-1]) + except FileExistsError: + # Multiple tensor-parallel ranks may prepare the same cache. + pass - current = [path for path in os.getenv("LIBRARY_PATH", "").split(":") if path] - os.environ["LIBRARY_PATH"] = ":".join( - dict.fromkeys([str(link_dir), str(library_dir), *current]) - ) - return [f"-Wl,-rpath,{library_dir}"] + return [f"-L{link_dir}", f"-Wl,-rpath,{library_dir}"] raise RuntimeError("Unable to locate libamdhip64 for ROCm JIT linking") diff --git a/tests/utils/test_rocm_arch.py b/tests/utils/test_rocm_arch.py index f080b1e93..955298623 100644 --- a/tests/utils/test_rocm_arch.py +++ b/tests/utils/test_rocm_arch.py @@ -1,3 +1,5 @@ +import importlib +import pathlib from types import SimpleNamespace import torch @@ -51,3 +53,37 @@ def test_hip_cflags_emit_one_offload_flag_per_arch(monkeypatch): assert "--offload-arch=gfx1200" in flags assert "--offload-arch=gfx1201" in flags assert not any(";" in flag for flag in flags) + + +def test_rocm_link_flags_support_versioned_modular_sdk(monkeypatch, tmp_path): + import torch.utils.cpp_extension as cpp_extension + + from freetoken.kernel import utils + + sdk = tmp_path / "sdk" + library_dir = sdk / "lib" + library_dir.mkdir(parents=True) + versioned_runtime = library_dir / "libamdhip64.so.7" + versioned_runtime.write_bytes(b"") + real_find_spec = importlib.util.find_spec + + def find_spec(name: str): + if name == "_rocm_sdk_core": + return SimpleNamespace(submodule_search_locations=[str(sdk)]) + return real_find_spec(name) + + monkeypatch.delenv("ROCM_HOME", raising=False) + monkeypatch.setattr(cpp_extension, "ROCM_HOME", None) + monkeypatch.setattr(importlib.util, "find_spec", find_spec) + monkeypatch.setattr(pathlib.Path, "home", lambda: tmp_path) + utils._rocm_link_flags.cache_clear() + + flags = utils._rocm_link_flags() + + compat_dir = tmp_path / ".cache" / "freetoken" / "rocm-lib" + compat_link = compat_dir / "libamdhip64.so" + assert f"-L{compat_dir}" in flags + assert f"-Wl,-rpath,{library_dir}" in flags + assert compat_link.resolve() == versioned_runtime.resolve() + + utils._rocm_link_flags.cache_clear() From c0713e44eace3cacd02f27f13d0f6032f4bc86ae Mon Sep 17 00:00:00 2001 From: zihaomu Date: Mon, 31 Aug 2026 15:52:32 +0800 Subject: [PATCH 09/10] test(rocm): harden foundation platform contracts --- .../csrc/include/freetoken/hip_compat.h | 4 +- .../kernel/csrc/include/freetoken/utils.cuh | 2 +- python/freetoken/utils/__init__.py | 4 -- python/freetoken/utils/arch.py | 19 ++----- tests/kernels/test_e4m3_compat.py | 23 +++++++++ tests/kernels/test_rocm_launch_kwargs.py | 51 +++++++++++++++++++ tests/utils/test_rocm_arch.py | 4 -- 7 files changed, 81 insertions(+), 26 deletions(-) create mode 100644 tests/kernels/test_rocm_launch_kwargs.py diff --git a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h index eb7e71e18..1898b4118 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h +++ b/python/freetoken/kernel/csrc/include/freetoken/hip_compat.h @@ -45,11 +45,11 @@ #endif #ifndef cudaMallocHost -#define cudaMallocHost hipMallocHost +#define cudaMallocHost hipHostMalloc #endif #ifndef cudaFreeHost -#define cudaFreeHost hipFreeHost +#define cudaFreeHost hipHostFree #endif #ifndef cudaHostAlloc diff --git a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh index f21b9585f..72a495969 100644 --- a/python/freetoken/kernel/csrc/include/freetoken/utils.cuh +++ b/python/freetoken/kernel/csrc/include/freetoken/utils.cuh @@ -116,7 +116,7 @@ public: } auto with_attr(bool use_pdl) -> LaunchKernel & { -#ifdef __HIP__ +#if FREETOKEN_USE_ROCM (void)use_pdl; m_config.numAttrs = 0; #else diff --git a/python/freetoken/utils/__init__.py b/python/freetoken/utils/__init__.py index c54579a16..af68bf581 100644 --- a/python/freetoken/utils/__init__.py +++ b/python/freetoken/utils/__init__.py @@ -2,8 +2,6 @@ is_arch_supported, is_rocm, get_rocm_gfx_arch, - is_gfx11xx_family, - is_gfx12xx_family, is_sm90_family, is_sm90_supported, is_sm100_family, @@ -41,8 +39,6 @@ "is_arch_supported", "is_rocm", "get_rocm_gfx_arch", - "is_gfx11xx_family", - "is_gfx12xx_family", "is_sm90_family", "is_sm90_supported", "is_sm100_family", diff --git a/python/freetoken/utils/arch.py b/python/freetoken/utils/arch.py index 3bf8fb61e..304505966 100644 --- a/python/freetoken/utils/arch.py +++ b/python/freetoken/utils/arch.py @@ -27,7 +27,10 @@ def get_rocm_gfx_arch() -> str | None: Prefer the runtime device because build variables may contain multiple semicolon-separated targets. Environment variables remain useful for - cross-compilation and systems where no GPU is currently visible. + cross-compilation and systems where no GPU is currently visible; in that + fallback mode the first target is returned. The result is process-cached + for FreeToken's one-process-per-GPU execution model, so callers must select + the intended device before the first call. """ if not is_rocm(): return None @@ -51,20 +54,6 @@ def get_rocm_gfx_arch() -> str | None: return None -@functools.cache -def is_gfx11xx_family() -> bool: - """True when the current AMD GPU is RDNA3 (gfx110x).""" - arch = get_rocm_gfx_arch() - return arch is not None and arch.startswith("gfx110") - - -@functools.cache -def is_gfx12xx_family() -> bool: - """True when the current AMD GPU is RDNA4 (gfx120x).""" - arch = get_rocm_gfx_arch() - return arch is not None and arch.startswith("gfx120") - - @functools.cache def _get_torch_cuda_version() -> Tuple[int, int] | None: import torch diff --git a/tests/kernels/test_e4m3_compat.py b/tests/kernels/test_e4m3_compat.py index 8bcb8e581..fa8a1f1d3 100644 --- a/tests/kernels/test_e4m3_compat.py +++ b/tests/kernels/test_e4m3_compat.py @@ -41,6 +41,29 @@ def _native_cc() -> bool: return torch.cuda.get_device_capability() >= (8, 9) +@pytest.mark.skipif(torch.version.hip is None, reason="needs ROCm") +def test_rocm_host_and_compile_time_native_gates_match(): + """The host wrapper and Triton constexpr must choose the same buffer ABI.""" + import triton + import triton.language as tl + + from freetoken.kernel.triton.e4m3_compat import e4m3_native, e4m3_native_cx + + @triton.jit + def gate_kernel(output): + if e4m3_native_cx(): + value = 1 + else: + value = 0 + tl.store(output, value) + + output = torch.empty(1, dtype=torch.int32, device="cuda") + gate_kernel[(1,)](output) + + assert e4m3_native() is False + assert output.item() == int(e4m3_native()) + + # ====================================================================================== # 1. Primitives vs the native fp8 unit (needs sm_89+ hardware for the reference). # ====================================================================================== diff --git a/tests/kernels/test_rocm_launch_kwargs.py b/tests/kernels/test_rocm_launch_kwargs.py new file mode 100644 index 000000000..52fac5451 --- /dev/null +++ b/tests/kernels/test_rocm_launch_kwargs.py @@ -0,0 +1,51 @@ +import pytest +import torch + + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available() or torch.version.hip is None, + reason="a ROCm GPU is required", +) + + +class _KernelRecorder: + def __init__(self): + self.kwargs = None + + def __getitem__(self, _grid): + def launch(*_args, **kwargs): + self.kwargs = kwargs + + return launch + + +def test_rocm_activation_launch_omits_cuda_pdl_attribute(monkeypatch): + import freetoken.kernel.triton.activation as activation + + recorder = _KernelRecorder() + monkeypatch.setattr(activation, "_act_and_mul_kernel", recorder) + monkeypatch.setattr(activation, "_pdl_supported", lambda: True) + monkeypatch.setattr(activation, "is_rocm", lambda: True) + + activation.silu_and_mul(torch.ones((2, 128), device="cuda")) + + assert recorder.kwargs is not None + assert recorder.kwargs["ENABLE_PDL"] is True + assert "launch_pdl" not in recorder.kwargs + + +def test_rocm_norm_launch_omits_cuda_pdl_attribute(monkeypatch): + import freetoken.kernel.triton.norm as norm + + recorder = _KernelRecorder() + monkeypatch.setattr(norm, "_rmsnorm_kernel", recorder) + monkeypatch.setattr(norm, "is_sm90_supported", lambda: True) + monkeypatch.setattr(norm, "is_rocm", lambda: True) + + x = torch.ones((2, 64), device="cuda") + weight = torch.ones(64, device="cuda") + norm.rmsnorm(x, weight) + + assert recorder.kwargs is not None + assert recorder.kwargs["ENABLE_PDL"] is True + assert "launch_pdl" not in recorder.kwargs diff --git a/tests/utils/test_rocm_arch.py b/tests/utils/test_rocm_arch.py index 955298623..224a61647 100644 --- a/tests/utils/test_rocm_arch.py +++ b/tests/utils/test_rocm_arch.py @@ -9,8 +9,6 @@ def _clear_arch_caches() -> None: arch.get_rocm_gfx_arch.cache_clear() - arch.is_gfx11xx_family.cache_clear() - arch.is_gfx12xx_family.cache_clear() def test_rocm_arch_prefers_visible_device_over_multi_arch_build_env(monkeypatch): @@ -26,8 +24,6 @@ def test_rocm_arch_prefers_visible_device_over_multi_arch_build_env(monkeypatch) _clear_arch_caches() assert arch.get_rocm_gfx_arch() == "gfx1201" - assert arch.is_gfx12xx_family() - assert not arch.is_gfx11xx_family() _clear_arch_caches() From e67f08eacbfa1403d5e2d1509e681f572ffa84a7 Mon Sep 17 00:00:00 2001 From: zihaomu Date: Fri, 4 Sep 2026 15:32:03 +0800 Subject: [PATCH 10/10] fix(rocm): make CPU MoE graph replay safe Fail closed to eager execution when the HIP stream-memory handshake cannot survive capture and replay. Add a ROCm 7.14 graph batch-memop path with executor-owned signal and parameter storage, dynamically size graph flag slots, preserve the existing CUDA module API, and cover the safety and multi-format replay paths. --- docs/models.md | 13 + python/freetoken/engine/engine.py | 94 +++- .../kernel/csrc/cpu_moe/cpu_moe_ext.cpp | 505 ++++++++++++++++-- python/freetoken/moe/cpu_executor.py | 183 +++++-- tests/engine/test_cpu_moe_graph_safety.py | 155 ++++++ tests/moe/test_cpu_moe.py | 131 ++++- tests/moe/test_cpu_moe_q4_0.py | 6 +- 7 files changed, 980 insertions(+), 107 deletions(-) create mode 100644 tests/engine/test_cpu_moe_graph_safety.py diff --git a/docs/models.md b/docs/models.md index e4850a124..16c5c8358 100644 --- a/docs/models.md +++ b/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 diff --git a/python/freetoken/engine/engine.py b/python/freetoken/engine/engine.py index 22c4a6c7e..20ac6eb4d 100644 --- a/python/freetoken/engine/engine.py +++ b/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 " 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 56ab93df8..f0f0f0550 100644 --- a/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp +++ b/python/freetoken/kernel/csrc/cpu_moe/cpu_moe_ext.cpp @@ -16,11 +16,12 @@ // (AVX-512-BF16 dpbf16 -> AVX-512F widening -> AVX2+FMA -> scalar). #include +#include #include -#include +#include #include +#include #include -#include #include #include #include @@ -558,15 +559,237 @@ float dot_nvfp4_i8_avx512vnni(const uint8_t* packed, const uint8_t* scale, float #endif // ===================================================================================== -// CUDA stream memory operations (driver API, resolved via dlopen -- no link-time or -// toolchain dependence). The GPU side of the flag handshake: submit = WRITE_VALUE -// (done[slot]=0 then ready[slot]=1), sync = WAIT_VALUE(done[slot] >= 1). The wait is -// executed by the GPU front-end (no SM-resident kernel), so GPU "utilization" stays -// truthful during CPU compute windows -- a resident spin kernel pinned it at 99%, -// which laptop CPU/GPU dynamic power schedulers answered by clamping the CPU's max -// frequency (GEMV workers -1.5x: the reported edge regression). Availability is -// probed functionally at startup (memops_probe); anything unsupported (Windows WDDM, -// vGPU, old drivers) falls back to the cudaLaunchHostFunc path. +// CUDA/HIP stream memory operations. The GPU side of the flag handshake is: +// submit = WRITE_VALUE(done=0), WRITE_VALUE(ready=slot+1); sync = WAIT_VALUE(done>=1). +// The wait runs on the GPU front end rather than occupying an SM. CUDA uses mapped +// pinned arrays and resolves the driver entry points at runtime. ROCm uses two shared +// 64-bit signals allocated with hipMallocSignalMemory (HIP requires that allocation +// type for WAIT addresses) and writes the task slot into the single ready signal. +// +// HIP's ordinary hipStreamWrite/WaitValue and hipStreamBatchMemOp calls execute while +// a stream is being captured but, as of ROCm 7.14, are not themselves recorded in the +// graph. During capture we therefore add explicit batch-memory-op graph nodes and splice +// them into the stream's current dependency set. The ROCm capability probe below does +// a real capture + instantiate + replay with a host handshake; eager enqueue alone is +// deliberately insufficient to enable the path. + +#if FREETOKEN_USE_ROCM && defined(HIP_VERSION_MAJOR) && defined(HIP_VERSION_MINOR) && \ + ((HIP_VERSION_MAJOR > 7) || (HIP_VERSION_MAJOR == 7 && HIP_VERSION_MINOR >= 14)) +#define FREETOKEN_HAS_HIP_GRAPH_MEMOPS 1 +#else +#define FREETOKEN_HAS_HIP_GRAPH_MEMOPS 0 +#endif + +#if FREETOKEN_USE_ROCM + +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS +static hipError_t hip_add_capture_memop_node( + hipStream_t stream, hipStreamBatchMemOpParams* ops, unsigned int count) { + if (ops == nullptr || count == 0) return hipErrorInvalidValue; + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipGraph_t graph = nullptr; + const hipGraphNode_t* dependencies = nullptr; + size_t dependency_count = 0; + hipError_t rc = hipStreamGetCaptureInfo_v2( + stream, &status, nullptr, &graph, &dependencies, &dependency_count); + if (rc != hipSuccess) return rc; + if (status != hipStreamCaptureStatusActive || graph == nullptr) + return hipErrorIllegalState; + + hipCtx_t context = nullptr; +#if defined(__GNUC__) +#pragma GCC diagnostic push +#pragma GCC diagnostic ignored "-Wdeprecated-declarations" +#endif + rc = hipCtxGetCurrent(&context); +#if defined(__GNUC__) +#pragma GCC diagnostic pop +#endif + if (rc != hipSuccess) return rc; + + hipBatchMemOpNodeParams params{}; + params.ctx = context; + params.count = count; + params.paramArray = ops; + hipGraphNode_t node = nullptr; + rc = hipGraphAddBatchMemOpNode( + &node, graph, dependencies, dependency_count, ¶ms); + if (rc != hipSuccess) return rc; + return hipStreamUpdateCaptureDependencies( + stream, &node, 1, hipStreamSetCaptureDependencies); +} + +static hipError_t hip_alloc_signal(uint64_t** out) { + void* storage = nullptr; + hipError_t rc = hipExtMallocWithFlags( + &storage, sizeof(uint64_t), hipMallocSignalMemory); + if (rc == hipSuccess) { + *out = static_cast(storage); + __atomic_store_n(*out, 0ULL, __ATOMIC_RELEASE); + } + return rc; +} + +static void hip_free_signal(uint64_t* signal) { + if (signal != nullptr) (void)hipFree(signal); +} + +static bool hipmemops_probe(uintptr_t /*stream*/, uintptr_t /*scratch_addr*/) { + int device = 0; + int can_wait = 0; + if (hipGetDevice(&device) != hipSuccess || + hipDeviceGetAttribute(&can_wait, hipDeviceAttributeCanUseStreamWaitValue, device) != + hipSuccess || + can_wait == 0) + return false; + + uint64_t* ready = nullptr; + uint64_t* done = nullptr; + hipStream_t probe_stream = nullptr; + hipGraph_t graph = nullptr; + hipGraphExec_t graph_exec = nullptr; + hipError_t rc = hip_alloc_signal(&ready); + if (rc == hipSuccess) rc = hip_alloc_signal(&done); + if (rc == hipSuccess) rc = hipStreamCreateWithFlags(&probe_stream, hipStreamNonBlocking); + + hipStreamBatchMemOpParams submit_ops[2]{}; + submit_ops[0].writeValue.operation = hipStreamMemOpWriteValue64; + submit_ops[0].writeValue.address = reinterpret_cast(done); + submit_ops[0].writeValue.value64 = 0; + submit_ops[1].writeValue.operation = hipStreamMemOpWriteValue64; + submit_ops[1].writeValue.address = reinterpret_cast(ready); + submit_ops[1].writeValue.value64 = 1; + hipStreamBatchMemOpParams sync_op{}; + sync_op.waitValue.operation = hipStreamMemOpWaitValue64; + sync_op.waitValue.address = reinterpret_cast(done); + sync_op.waitValue.value64 = 1; + sync_op.waitValue.flags = hipStreamWaitValueGte; + + if (rc == hipSuccess) + rc = hipStreamBeginCapture(probe_stream, hipStreamCaptureModeRelaxed); + if (rc == hipSuccess) + rc = hip_add_capture_memop_node(probe_stream, submit_ops, 2); + if (rc == hipSuccess) + rc = hip_add_capture_memop_node(probe_stream, &sync_op, 1); + if (rc == hipSuccess) rc = hipStreamEndCapture(probe_stream, &graph); + if (rc == hipSuccess) + rc = hipGraphInstantiate(&graph_exec, graph, nullptr, nullptr, 0); + + bool served = false; + std::thread responder; + if (rc == hipSuccess) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + __atomic_store_n(done, 0ULL, __ATOMIC_RELEASE); + responder = std::thread([&] { + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(2); + while (std::chrono::steady_clock::now() < deadline) { + if (__atomic_load_n(ready, __ATOMIC_ACQUIRE) == 1ULL) { + __atomic_store_n(ready, 0ULL, __ATOMIC_RELEASE); + served = true; + break; + } + std::this_thread::yield(); + } + // Always release the graph wait, including on probe failure, so startup cannot + // wedge indefinitely on a beta runtime implementation. + __atomic_store_n(done, 1ULL, __ATOMIC_RELEASE); + }); + rc = hipGraphLaunch(graph_exec, probe_stream); + if (rc == hipSuccess) rc = hipStreamSynchronize(probe_stream); + } + if (responder.joinable()) responder.join(); + const bool ok = rc == hipSuccess && served && + __atomic_load_n(done, __ATOMIC_ACQUIRE) == 1ULL; + + if (graph_exec != nullptr) (void)hipGraphExecDestroy(graph_exec); + if (graph != nullptr) (void)hipGraphDestroy(graph); + if (probe_stream != nullptr) (void)hipStreamDestroy(probe_stream); + hip_free_signal(ready); + hip_free_signal(done); + return ok; +} + +static void hipmemop_throw(hipError_t rc, const char* what) { + if (rc != hipSuccess) { + throw std::runtime_error( + std::string("CPU MoE HIP flag handshake: ") + what + " failed: " + + hipGetErrorString(rc)); + } +} + +static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, + int64_t slot, uintptr_t capture_ops_addr = 0) { + auto hip_stream = reinterpret_cast(stream); + auto* done = reinterpret_cast(done_addr); + auto* ready = reinterpret_cast(ready_addr); + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipError_t rc = hipStreamIsCapturing(hip_stream, &status); + hipmemop_throw(rc, "submit capture-status query"); + if (status != hipStreamCaptureStatusActive) { + hipmemop_throw( + hipStreamWriteValue64(hip_stream, done, 0, hipStreamWriteValueDefault), + "hipStreamWriteValue64(done)"); + hipmemop_throw( + hipStreamWriteValue64( + hip_stream, ready, static_cast(slot + 1), + hipStreamWriteValueDefault), + "hipStreamWriteValue64(ready)"); + return; + } + + auto* capture_ops = reinterpret_cast(capture_ops_addr); + if (capture_ops == nullptr) { + throw std::runtime_error( + "CPU MoE HIP graph memops require executor-owned parameter storage"); + } + hipmemop_throw(hip_add_capture_memop_node(hip_stream, capture_ops, 2), + "submit graph batch-memory-op"); +} + +static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t /*slot*/, + uintptr_t capture_ops_addr = 0) { + auto hip_stream = reinterpret_cast(stream); + auto* done = reinterpret_cast(done_addr); + hipStreamCaptureStatus status = hipStreamCaptureStatusNone; + hipError_t rc = hipStreamIsCapturing(hip_stream, &status); + hipmemop_throw(rc, "sync capture-status query"); + if (status != hipStreamCaptureStatusActive) { + hipmemop_throw( + hipStreamWaitValue64( + hip_stream, done, 1, hipStreamWaitValueGte, UINT64_MAX), + "hipStreamWaitValue64(done)"); + return; + } + + auto* capture_ops = reinterpret_cast(capture_ops_addr); + if (capture_ops == nullptr) { + throw std::runtime_error( + "CPU MoE HIP graph memops require executor-owned parameter storage"); + } + hipmemop_throw(hip_add_capture_memop_node(hip_stream, capture_ops, 1), + "sync graph batch-memory-op"); +} + +#else // ROCm headers older than the graph batch-memory-op API + +static bool hipmemops_probe(uintptr_t, uintptr_t) { return false; } +static void cumemop_submit(uintptr_t, uintptr_t, uintptr_t, int64_t, uintptr_t = 0) { + throw std::runtime_error( + "CPU MoE HIP flag handshake requires ROCm 7.14 graph batch-memory-op APIs"); +} +static void cumemop_sync(uintptr_t, uintptr_t, int64_t, uintptr_t = 0) { + throw std::runtime_error( + "CPU MoE HIP flag handshake requires ROCm 7.14 graph batch-memory-op APIs"); +} + +#endif // FREETOKEN_HAS_HIP_GRAPH_MEMOPS + +static bool cumemops_probe(uintptr_t stream, uintptr_t scratch_addr) { + return hipmemops_probe(stream, scratch_addr); +} + +#else // CUDA + #if defined(_WIN32) #include static void* cumemop_dlopen() { return (void*)::LoadLibraryA("nvcuda.dll"); } @@ -636,7 +859,7 @@ static void cumemop_check(int rc, const char* what) { } static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, - int64_t slot) { + int64_t slot, uintptr_t /*capture_ops_addr*/ = 0) { auto* s = reinterpret_cast(stream); // Order matters and is preserved by the front end: reset done BEFORE raising ready, // so the coordinator's completion write for THIS step can never be wiped. @@ -648,13 +871,16 @@ static void cumemop_submit(uintptr_t stream, uintptr_t done_addr, uintptr_t read "cuStreamWriteValue64(ready)"); } -static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot) { +static void cumemop_sync(uintptr_t stream, uintptr_t done_addr, int64_t slot, + uintptr_t /*capture_ops_addr*/ = 0) { cumemop_check(g_cu_wait64(reinterpret_cast(stream), (unsigned long long)(done_addr + (size_t)slot * 8), 1ULL, kCuWaitValueGeq), "cuStreamWaitValue64(done)"); } +#endif // FREETOKEN_USE_ROCM + struct DotChoice { dot_fn fn; const char* name; @@ -1301,17 +1527,28 @@ struct CpuMoeExecutor { std::vector core_ids; // worker tid -> logical CPU to pin to (may be empty) // ---- Flag-based GPU<->CPU handshake (replaces the per-layer cudaLaunchHostFunc pair) ---- - // A tiny GPU kernel bumps ready_flags[slot] at submit; this coordinator thread busy-polls - // it, runs the slot's task on the worker pool, and sets done_flags[slot], which a GPU - // spin-wait kernel polls at sync. This removes the ~2x30-50us host-func dispatch round - // trips per MoE layer per decode step that otherwise idle the GPU (~6 ms/step on a - // 75-layer model). One slot per (layer, decode batch size) pair -- the Python side - // allocates slots as tasks are created. Flags live in mapped-pinned host memory (UVA: - // the same pointers are used by the GPU kernels and by this thread). + // CUDA uses one mapped-pinned ready/done entry per task slot. HIP signal memory is + // restricted to one 64-bit word per allocation, so ROCm uses a shared ready command + // (slot+1) and shared done flag. Decode layers are serialized by their stream waits, + // therefore only one CPU task can own that pair at a time. std::thread coord_thread; std::atomic coord_stop{false}; - volatile int64_t* ready_flags = nullptr; // GPU increments, this thread polls - volatile int64_t* done_flags = nullptr; // this thread sets, GPU spin-waits + volatile int64_t* ready_flags = nullptr; + volatile int64_t* done_flags = nullptr; + bool shared_flag_signal = false; +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + uint64_t* owned_ready_signal = nullptr; + uint64_t* owned_done_signal = nullptr; + // ROCm 7.14's beta graph API retains hipBatchMemOpNodeParams::paramArray. + // There are only two immutable node shapes per task slot (submit has two ops; + // sync has one shared op), so allocate one canonical array for each shape once + // and let every rebuilt graph reference it. The vector is sized before capture + // starts and never resized, keeping addresses stable without a module-global, + // append-only allocation. These arrays and the signals now share the executor's + // lifetime, which is also the minimum lifetime required by captured CPU-MoE nodes. + std::vector> hip_graph_submit_ops; + std::array hip_graph_sync_ops{}; +#endif int coord_num_slots = 0; std::vector flag_task; // slot -> task (registered lazily) std::vector flag_served; // slot -> completed dispatch count (tests/debug) @@ -1532,6 +1769,12 @@ struct CpuMoeExecutor { ~CpuMoeExecutor() { coord_stop.store(true); if (coord_thread.joinable()) coord_thread.join(); +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + hip_free_signal(owned_ready_signal); + hip_free_signal(owned_done_signal); + owned_ready_signal = nullptr; + owned_done_signal = nullptr; +#endif { std::lock_guard lk(task_mtx); stop = true; @@ -1979,11 +2222,72 @@ struct CpuMoeExecutor { return (slot >= 0 && slot < static_cast(flag_served.size())) ? flag_served[slot] : 0; } + uintptr_t flag_ready_address() const { + return reinterpret_cast(ready_flags); + } + + uintptr_t flag_done_address() const { + return reinterpret_cast(done_flags); + } + + int64_t flag_ready_value(int slot) const { + if (ready_flags == nullptr) return 0; + if (shared_flag_signal) { + const int64_t command = flag_load_acquire(&ready_flags[0]); + return command == slot + 1 ? 1 : 0; + } + return flag_load_acquire(&ready_flags[slot]); + } + + int64_t flag_done_value(int slot) const { + if (done_flags == nullptr) return 0; + return flag_load_acquire(&done_flags[shared_flag_signal ? 0 : slot]); + } + + std::vector pending_flag_slots() const { + std::vector pending; + if (ready_flags == nullptr || done_flags == nullptr) return pending; + if (shared_flag_signal) { + const int64_t command = flag_load_acquire(&ready_flags[0]); + if (command > 0 && command <= coord_num_slots && + flag_load_acquire(&done_flags[0]) == 0) + pending.push_back(static_cast(command - 1)); + return pending; + } + for (int slot = 0; slot < coord_num_slots; ++slot) { + if (flag_load_acquire(&ready_flags[slot]) != 0 && + flag_load_acquire(&done_flags[slot]) == 0) + pending.push_back(slot); + } + return pending; + } + + void poison_flag(int slot) { + if (done_flags != nullptr) + flag_store_release(&done_flags[shared_flag_signal ? 0 : slot], 1); + } + + void launch_flag_coordinator(int pin_core) { + coord_stop.store(false); + coord_thread = std::thread([this, pin_core] { +#if CPU_MOE_HAS_AFFINITY + if (pin_core >= 0) { + cpu_set_t set; + CPU_ZERO(&set); + CPU_SET(pin_core, &set); + pthread_setaffinity_np(pthread_self(), sizeof(set), &set); + } +#endif + coordinator_loop(); + }); + } + // Start the busy-poll coordinator over the mapped-pinned flag arrays. ``pin_core`` >= 0 // pins the coordinator to that logical CPU (the worker auto-sizing reserves it), so its // polling never migrates onto / contends with a GEMV worker's core. void start_flag_coordinator(uintptr_t ready_ptr, uintptr_t done_ptr, int num_slots, int pin_core) { + shared_flag_signal = false; ready_flags = reinterpret_cast(ready_ptr); done_flags = reinterpret_cast(done_ptr); coord_num_slots = num_slots; @@ -1992,18 +2296,102 @@ struct CpuMoeExecutor { if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); } flag_served.assign(num_slots, 0); - coord_stop.store(false); - coord_thread = std::thread([this, pin_core] { -#if CPU_MOE_HAS_AFFINITY - if (pin_core >= 0) { - cpu_set_t set; - CPU_ZERO(&set); - CPU_SET(pin_core, &set); - pthread_setaffinity_np(pthread_self(), sizeof(set), &set); - } + launch_flag_coordinator(pin_core); + } + + void start_shared_signal_coordinator(int num_slots, int pin_core) { +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + if (num_slots <= 0) + throw std::invalid_argument("CPU MoE HIP flag slot count must be positive"); + if (owned_ready_signal != nullptr || owned_done_signal != nullptr || + !hip_graph_submit_ops.empty()) + throw std::logic_error( + "CPU MoE HIP flag coordinator is already initialized"); + hipError_t rc = hip_alloc_signal(&owned_ready_signal); + if (rc == hipSuccess) rc = hip_alloc_signal(&owned_done_signal); + if (rc != hipSuccess) { + hip_free_signal(owned_ready_signal); + hip_free_signal(owned_done_signal); + owned_ready_signal = nullptr; + owned_done_signal = nullptr; + throw std::runtime_error( + std::string("CPU MoE HIP signal allocation failed: ") + hipGetErrorString(rc)); + } + shared_flag_signal = true; + ready_flags = reinterpret_cast(owned_ready_signal); + done_flags = reinterpret_cast(owned_done_signal); + coord_num_slots = num_slots; + hip_graph_submit_ops.clear(); + hip_graph_submit_ops.resize(num_slots); + for (int slot = 0; slot < num_slots; ++slot) { + auto& ops = hip_graph_submit_ops[slot]; + ops[0].writeValue.operation = hipStreamMemOpWriteValue64; + ops[0].writeValue.address = reinterpret_cast(owned_done_signal); + ops[0].writeValue.value64 = 0; + ops[1].writeValue.operation = hipStreamMemOpWriteValue64; + ops[1].writeValue.address = reinterpret_cast(owned_ready_signal); + ops[1].writeValue.value64 = static_cast(slot + 1); + } + hip_graph_sync_ops[0] = {}; + hip_graph_sync_ops[0].waitValue.operation = hipStreamMemOpWaitValue64; + hip_graph_sync_ops[0].waitValue.address = + reinterpret_cast(owned_done_signal); + hip_graph_sync_ops[0].waitValue.value64 = 1; + hip_graph_sync_ops[0].waitValue.flags = hipStreamWaitValueGte; + { + std::lock_guard lk(flag_task_mtx); + if (static_cast(flag_task.size()) < num_slots) flag_task.resize(num_slots, nullptr); + } + flag_served.assign(num_slots, 0); + launch_flag_coordinator(pin_core); +#else + (void)num_slots; + (void)pin_core; + throw std::runtime_error( + "CPU MoE HIP signal allocation requires ROCm 7.14 graph memops"); +#endif + } + + void submit_flag_memop(uintptr_t stream, uintptr_t done_addr, + uintptr_t ready_addr, int64_t slot) { + uintptr_t capture_ops_addr = 0; +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + if (shared_flag_signal) { + if (slot < 0 || slot >= static_cast(hip_graph_submit_ops.size())) + throw std::out_of_range("CPU MoE HIP flag slot is out of range"); + if (done_addr != reinterpret_cast(owned_done_signal) || + ready_addr != reinterpret_cast(owned_ready_signal)) + throw std::invalid_argument( + "CPU MoE HIP memops must use this executor's signal memory"); + capture_ops_addr = + reinterpret_cast(hip_graph_submit_ops[slot].data()); + } +#endif + cumemop_submit(stream, done_addr, ready_addr, slot, capture_ops_addr); + } + + void sync_flag_memop(uintptr_t stream, uintptr_t done_addr, int64_t slot) { + uintptr_t capture_ops_addr = 0; +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + if (shared_flag_signal) { + if (slot < 0 || slot >= static_cast(hip_graph_submit_ops.size())) + throw std::out_of_range("CPU MoE HIP flag slot is out of range"); + if (done_addr != reinterpret_cast(owned_done_signal)) + throw std::invalid_argument( + "CPU MoE HIP memops must use this executor's signal memory"); + capture_ops_addr = reinterpret_cast(hip_graph_sync_ops.data()); + } +#endif + cumemop_sync(stream, done_addr, slot, capture_ops_addr); + } + + size_t graph_memop_param_count() const { +#if FREETOKEN_HAS_HIP_GRAPH_MEMOPS + return hip_graph_submit_ops.size() * 2 + + (hip_graph_submit_ops.empty() ? 0 : hip_graph_sync_ops.size()); +#else + return 0; #endif - coordinator_loop(); - }); } void coordinator_loop() { @@ -2026,13 +2414,19 @@ struct CpuMoeExecutor { bool dozing = false; while (!coord_stop.load(std::memory_order_relaxed)) { bool any = false; - for (int L = 0; L < coord_num_slots; ++L) { + const int begin = shared_flag_signal + ? static_cast(flag_load_acquire(&ready_flags[0]) - 1) + : 0; + const int end = shared_flag_signal ? begin + 1 : coord_num_slots; + for (int L = begin; L < end; ++L) { + if (L < 0 || L >= coord_num_slots) continue; // Binary handshake (memop-compatible: the GPU-side WAIT compares against an // immediate baked at graph capture, so the protocol resets per step instead of // counting). Acquire: everything the GPU made visible before setting ready -- // the D2H input copies -- is visible to the worker pool after this read. - if (flag_load_acquire(&ready_flags[L]) != 0) { - flag_store_release(&ready_flags[L], 0); // consume this step's doorbell + volatile int64_t* ready = &ready_flags[shared_flag_signal ? 0 : L]; + if (flag_load_acquire(ready) != 0) { + flag_store_release(ready, 0); // consume this step's doorbell MoeTask* t; { std::lock_guard lk(flag_task_mtx); @@ -2043,7 +2437,7 @@ struct CpuMoeExecutor { sync(); } // Release: the workers' y stores are visible before the GPU sees done. - flag_store_release(&done_flags[L], 1); + flag_store_release(&done_flags[shared_flag_signal ? 0 : L], 1); if (L < static_cast(flag_served.size())) ++flag_served[L]; any = true; } @@ -2077,7 +2471,8 @@ struct CpuMoeExecutor { // caught mid-shutdown exits its sync kernel now instead of owning the watchdog // stall. Runs before the destructor's join() returns, while the flag arrays are // still alive on the Python side. - for (int L = 0; L < coord_num_slots; ++L) { + const int done_count = shared_flag_signal ? 1 : coord_num_slots; + for (int L = 0; L < done_count; ++L) { flag_store_release(&done_flags[L], INT64_MAX); } } @@ -2128,18 +2523,46 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { .def("register_flag_task", &CpuMoeExecutor::register_flag_task, py::arg("slot"), py::arg("task")) .def("flag_served_count", &CpuMoeExecutor::flag_served_count, py::arg("slot")) + .def("flag_ready_address", &CpuMoeExecutor::flag_ready_address) + .def("flag_done_address", &CpuMoeExecutor::flag_done_address) + .def("flag_ready_value", &CpuMoeExecutor::flag_ready_value, py::arg("slot")) + .def("flag_done_value", &CpuMoeExecutor::flag_done_value, py::arg("slot")) + .def("pending_flag_slots", &CpuMoeExecutor::pending_flag_slots) + .def("poison_flag", &CpuMoeExecutor::poison_flag, py::arg("slot")) .def("start_flag_coordinator", &CpuMoeExecutor::start_flag_coordinator, py::arg("ready_ptr"), py::arg("done_ptr"), py::arg("num_slots"), py::arg("pin_core")) + .def("start_shared_signal_coordinator", + &CpuMoeExecutor::start_shared_signal_coordinator, + py::arg("num_slots"), py::arg("pin_core")) + .def("submit_flag_memop", &CpuMoeExecutor::submit_flag_memop, + py::arg("stream"), py::arg("done_addr"), py::arg("ready_addr"), + py::arg("slot")) + .def("sync_flag_memop", &CpuMoeExecutor::sync_flag_memop, + py::arg("stream"), py::arg("done_addr"), py::arg("slot")) + .def("graph_memop_param_count", &CpuMoeExecutor::graph_memop_param_count) .def("set_input_prequant", [](CpuMoeExecutor& e, bool v) { e.input_prequant = v; }, py::arg("value")) .def("isa_name", &CpuMoeExecutor::isa_name); m.def("memops_probe", &cumemops_probe, py::arg("stream"), py::arg("scratch_addr")); - m.def("memop_submit", &cumemop_submit, py::arg("stream"), py::arg("done_addr"), - py::arg("ready_addr"), py::arg("slot")); - m.def("memop_sync", &cumemop_sync, py::arg("stream"), py::arg("done_addr"), + // Preserve the established CUDA-facing module API. ROCm graph capture uses the + // executor methods above so its beta node parameter arrays have executor lifetime; + // eager ROCm calls remain valid here as well. + m.def("memop_submit", + [](uintptr_t stream, uintptr_t done_addr, uintptr_t ready_addr, int64_t slot) { + cumemop_submit(stream, done_addr, ready_addr, slot); + }, + py::arg("stream"), py::arg("done_addr"), py::arg("ready_addr"), py::arg("slot")); + m.def("memop_sync", + [](uintptr_t stream, uintptr_t done_addr, int64_t slot) { + cumemop_sync(stream, done_addr, slot); + }, + py::arg("stream"), py::arg("done_addr"), py::arg("slot")); + m.def("memops_use_shared_signal", []() { + return static_cast(FREETOKEN_HAS_HIP_GRAPH_MEMOPS); + }); // ABI capability marker: the highest ActKind this build implements in the // GENERIC epilogue. CpuMoeExecutor.__init__ probes it before requesting an act // id the epilogue must handle -- a prebuilt .so from before ACT_SWIGLUOAI diff --git a/python/freetoken/moe/cpu_executor.py b/python/freetoken/moe/cpu_executor.py index b96205aa6..9ed53d4db 100644 --- a/python/freetoken/moe/cpu_executor.py +++ b/python/freetoken/moe/cpu_executor.py @@ -1,18 +1,28 @@ """Python wrapper around the ``_cpu_moe`` C++ executor (--moe-backend cpu). Owns the persistent CPU worker pool, the per-batch-size pinned IO buffers, and -the per-(layer, batch-size) host-func task descriptors. ``decode`` issues the -whole CUDA-graph-capturable sequence on the current stream: +the per-(layer, batch-size) task descriptors. ``decode`` issues this sequence on +the current stream: D2H (hidden, topk_ids, topk_weights -> pinned) - -> submit host node (cudaLaunchHostFunc: enqueue MoE task to the pool) - -> sync host node (cudaLaunchHostFunc: spin until the pool drains) + -> submit flag memop (or eager cudaLaunchHostFunc fallback) + -> sync flag memop (or eager cudaLaunchHostFunc fallback) -> H2D (pinned expert output -> GPU) -Buffers and tasks are allocated lazily per batch size. GraphRunner runs an eager -``model.forward()`` at each batch size immediately before capturing it, so the -first (eager) call materializes the stable pinned buffers + task pointers that -the subsequent capture embeds in its host/memcpy nodes. +CUDA and ROCm expose the same ordering and results; only the GPU/CPU handshake +implementation differs. CUDA keeps mapped-pinned ready/done arrays and the +established module-level ``cuStreamWriteValue64``/``cuStreamWaitValue64`` path. +ROCm 7.14 uses executor-owned ``hipMallocSignalMemory`` signals and explicitly +adds batch-memory-op nodes during graph capture. The executor also owns the HIP +node-parameter arrays because captured graphs retain their addresses through +replay. + +Only a ROCm handshake that passes a real capture/instantiate/replay probe is +graph-safe; the Engine disables graph capture and continues eagerly when the +probe fails. Buffers and tasks are allocated lazily per batch size. GraphRunner +runs an eager ``model.forward()`` at each batch size immediately before +capturing it, so the first (eager) call materializes the stable pinned buffers + +task pointers that the subsequent capture embeds in its host/memcpy nodes. """ from __future__ import annotations @@ -32,26 +42,27 @@ # Flag-based GPU<->CPU handshake for hybrid/cpu decode. The default host-func path # (cudaLaunchHostFunc submit+sync per layer) pays ~30-50us of callback dispatch latency # per call with the GPU stream idle -- 2 calls per MoE layer per decode step (~6 ms/step -# on a 75-layer model). Instead the GPU raises a mapped-pinned "ready" flag at submit; a -# persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets a "done" -# flag the GPU waits on at sync -- no host-func round-trip. Both GPU-side operations are -# STREAM MEMORY OPERATIONS (cuStreamWriteValue64 / cuStreamWaitValue64, resolved from the -# driver at runtime): they execute on the GPU front end with no SM-resident kernel, so -# GPU utilization stays truthful during the CPU compute window. (The first cut used a -# spin-wait kernel; that pinned reported utilization at 99% and laptop CPU/GPU dynamic -# power schedulers responded by clamping the CPU frequency -- a net decode regression on -# power-coupled edge devices.) Each (layer, decode batch size) pair gets its own flag -# slot, so every captured decode graph rides the handshake. Where memops are unavailable -# (Windows WDDM, vGPU, old drivers -- functionally probed at startup) or the slot -# capacity is exceeded, decode keeps the host-func path (functional, slower). A Python -# watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] + +# on a 75-layer model). Instead the GPU raises a backend-specific "ready" signal at +# submit; a persistent CPU coordinator (in _cpu_moe) polls it, runs the layer, and sets +# the "done" signal the GPU waits on at sync -- no host-func round-trip. CUDA uses +# mapped-pinned per-slot arrays with cuStreamWriteValue64/cuStreamWaitValue64; ROCm uses +# executor-owned HIP signal memory and explicit graph batch-memory-op nodes. Both execute +# on the GPU front end with no SM-resident kernel, so GPU utilization stays truthful +# during the CPU compute window. (The first cut used a spin-wait kernel; that pinned +# reported utilization at 99% and laptop CPU/GPU dynamic power schedulers responded by +# clamping the CPU frequency -- a net decode regression on power-coupled edge devices.) +# Each (layer, decode batch size) pair gets a logical slot, so every captured decode graph +# rides the handshake. Where memops are unavailable (Windows WDDM, vGPU, old drivers -- +# functionally probed at startup), CUDA keeps the host-func path (functional, slower); +# ROCm fails closed during capture and the Engine normally selects the eager path. A +# Python watchdog thread turns a wedged coordinator into a loud RuntimeError (via err[] + # raise_if_unhealthy) instead of an indefinite stream stall. # Caveat: the coordinator busy-polls one core while decode traffic flows (idle backoff # otherwise); FREETOKEN_CPU_MOE_FLAG_SYNC=0 opts out entirely. _FLAG_SYNC = os.getenv("FREETOKEN_CPU_MOE_FLAG_SYNC", "1") != "0" -# Flag slots per MoE layer: covers this many distinct decode batch sizes (captured graph -# sizes plus any eager padded sizes); more than that is unheard of, and the overflow -# just keeps the host-func path for the extra combos. +_IS_ROCM = getattr(torch.version, "hip", None) is not None +# Minimum flag slots per MoE layer, retaining eager headroom for small graph sets. The +# Engine raises this to the number of distinct configured graph batch sizes when needed. _FLAG_SLOTS_PER_LAYER = 16 # Activation ids must match ActKind in csrc/cpu_moe/cpu_moe_ext.cpp. Id 3 is the @@ -154,6 +165,7 @@ def __init__( num_threads: int, max_tokens: int, device: torch.device, + flag_slots_per_layer: int = _FLAG_SLOTS_PER_LAYER, swiglu_alpha: float = 1.702, swiglu_limit: float | None = None, ) -> None: @@ -188,6 +200,9 @@ def __init__( self.quant_format = fmt self.device = device self.max_tokens = int(max_tokens) + self._flag_slots_per_layer = int(flag_slots_per_layer) + if self._flag_slots_per_layer < 1: + raise ValueError("flag_slots_per_layer must be positive") self.apply_router_weight_on_input = bool(apply_router_weight_on_input) # The per-layer tensors and their pointer tables must outlive the executor # (C++ holds raw addresses into both). @@ -195,7 +210,9 @@ def __init__( ptrs, (self.H, self.I) = self._resolve_banks(cache.bank_sources, fmt) # Decide the flag handshake up front (env + device + a functional stream-memop - # probe): its coordinator needs a core of its own, which the auto thread sizing + # probe). On ROCm the extension's probe performs a real graph capture + replay, + # not merely an eager enqueue: its coordinator needs a core of its own, which + # the auto thread sizing # below reserves (a coordinator time-slicing against the GEMV workers measurably # destabilizes throughput on fully-subscribed boxes). self._flag_sync = _FLAG_SYNC and device.type == "cuda" @@ -206,11 +223,18 @@ def __init__( if not _cpu_moe.memops_probe( torch.cuda.current_stream().cuda_stream, probe_scratch.data_ptr() ): - logger.info_rank0( - "cpu-moe flag handshake unavailable: CUDA stream memory operations " - "are not supported here (Windows WDDM / vGPU / old driver); using " - "the cudaLaunchHostFunc sync" - ) + if _IS_ROCM: + logger.info_rank0( + "cpu-moe flag handshake unavailable: stream memory operations " + "did not pass the capture/replay probe; using " + "cudaLaunchHostFunc only for eager execution" + ) + else: + logger.info_rank0( + "cpu-moe flag handshake unavailable: CUDA stream memory " + "operations are not supported here (Windows WDDM / vGPU / old " + "driver); using the cudaLaunchHostFunc sync" + ) self._flag_sync = False nthreads, core_ids = resolve_threads_and_affinity(num_threads) @@ -266,19 +290,31 @@ def __init__( # self so the coordinator's pinned pointers stay valid for the executor's # lifetime (flag_sync itself was decided above, before thread sizing). self._ready = self._done = self._err = None + self._shared_flag_signal = bool( + self._flag_sync + and getattr(_cpu_moe, "memops_use_shared_signal", lambda: False)() + ) self._flag_slots: dict[tuple[int, int], int] = {} # (layer_id, bs) -> slot - self._flag_capacity = self.num_layers * _FLAG_SLOTS_PER_LAYER + self._flag_capacity = self.num_layers * self._flag_slots_per_layer if self._flag_sync: - self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) self._err = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) - self._ready.zero_() - self._done.zero_() self._err.zero_() - self._ext.start_flag_coordinator( - self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity, - self._coord_core, - ) + if self._shared_flag_signal: + # HIP wait addresses must come from hipMallocSignalMemory. The C++ + # executor owns the two exact 64-bit allocations so it can stop/join + # the coordinator before freeing them. + self._ext.start_shared_signal_coordinator( + self._flag_capacity, self._coord_core + ) + else: + self._ready = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) + self._done = alloc_pinned_tensor(self._flag_capacity, dtype=torch.int64) + self._ready.zero_() + self._done.zero_() + self._ext.start_flag_coordinator( + self._ready.data_ptr(), self._done.data_ptr(), self._flag_capacity, + self._coord_core, + ) 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 @@ -317,6 +353,11 @@ def __init__( f"top_k={self.top_k} act={activation} max_tokens={self.max_tokens}" ) + @property + def graph_capture_safe(self) -> bool: + """Whether this executor passed the native stream-memop graph replay probe.""" + return self._flag_sync + def _make_table(self, layers: list[torch.Tensor]) -> torch.Tensor: """Build a CPU int64 tensor of per-layer base addresses for one bank. @@ -550,8 +591,33 @@ def decode_submit( computes only the routes assigned to it. Returns an opaque handle to pass to :meth:`decode_sync`. The output tensor is allocated here so it stays live (and distinct from the interleaved GPU work) across the overlap window.""" + # hipLaunchHostFunc is not replay-safe for this CPU handshake on ROCm. The + # Engine normally disables graphs when the native capture/replay probe failed; + # this is the non-bypassable last line of defence for direct executor use or a + # future configuration path that misses that downgrade. + if ( + _IS_ROCM + and not self._flag_sync + and torch.cuda.is_current_stream_capturing() + ): + raise RuntimeError( + "CPU/Hybrid MoE is not CUDA-graph safe on ROCm when " + "stream-memory synchronization is unavailable" + ) + bs = hidden_states.shape[0] io = self._io_for(bs) + task = self._task_for(layer_id, bs) + slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None + if ( + _IS_ROCM + and torch.cuda.is_current_stream_capturing() + and slot is None + ): + raise RuntimeError( + "CPU/Hybrid MoE exhausted its graph-safe ROCm flag slots during " + "CUDA-graph capture" + ) if self._gpu_prequant: # DSV4: apply the reference FP8 round-trip on the GPU (the same kernel the @@ -566,16 +632,24 @@ def decode_submit( io["ids"].copy_(topk_ids.to(torch.int32), non_blocking=True) io["w"].copy_(topk_weights.to(torch.float32), non_blocking=True) - task = self._task_for(layer_id, bs) out = torch.empty_like(hidden_states) - slot = self._flag_slots.get((layer_id, bs)) if self._flag_sync else None if slot is not None: # Front-end memops: done[slot]=0 then ready[slot]=1 (the coordinator's # doorbell). No kernel launched; no host-func round trip. - self._cpu_moe.memop_submit( - torch.cuda.current_stream().cuda_stream, - self._done.data_ptr(), self._ready.data_ptr(), slot, - ) + stream = torch.cuda.current_stream().cuda_stream + if self._shared_flag_signal: + self._ext.submit_flag_memop( + stream, + self._ext.flag_done_address(), + self._ext.flag_ready_address(), + slot, + ) + else: + # Keep CUDA on the established module-level driver-API path. Only + # ROCm needs executor-owned graph parameter storage. + self._cpu_moe.memop_submit( + stream, self._done.data_ptr(), self._ready.data_ptr(), slot + ) else: stream = torch.cuda.current_stream().cuda_stream self._ext.submit_with_cuda_stream(stream, task) @@ -589,9 +663,13 @@ def decode_sync(self, pending: tuple) -> torch.Tensor: if slot is not None: # Front-end WAIT(done[slot] >= 1): blocks this stream's later nodes without # occupying an SM, so GPU utilization stays truthful during the CPU window. - self._cpu_moe.memop_sync( - torch.cuda.current_stream().cuda_stream, self._done.data_ptr(), slot, - ) + stream = torch.cuda.current_stream().cuda_stream + if self._shared_flag_signal: + self._ext.sync_flag_memop( + stream, self._ext.flag_done_address(), slot + ) + else: + self._cpu_moe.memop_sync(stream, self._done.data_ptr(), slot) else: stream = torch.cuda.current_stream().cuda_stream self._ext.sync_with_cuda_stream(stream, task) @@ -610,12 +688,11 @@ def _watchdog_tick(self, suspects: dict) -> None: windows under heavy external load, but a coordinator that made progress in between is alive by definition. ``suspects`` maps slot -> (first_seen, served_at_first_sight) and persists across ticks.""" - stuck = (self._ready == 1) & (self._done == 0) - if not bool(stuck.any()): + pending = set(self._ext.pending_flag_slots()) + if not pending: suspects.clear() return now = time.monotonic() - pending = set(stuck.nonzero().flatten().tolist()) for slot in list(suspects): if slot not in pending: del suspects[slot] @@ -637,7 +714,7 @@ def _watchdog_tick(self, suspects: dict) -> None: for i in dead: self._err[i] = 1 for i in dead: - self._done[i] = 1 # after err: unblock the stream into a checked failure + self._ext.poison_flag(i) # after err: unblock into a checked failure suspects.pop(i, None) def raise_if_unhealthy(self) -> None: @@ -662,7 +739,7 @@ def _watchdog_main(executor_ref) -> None: while True: time.sleep(2.0) executor = executor_ref() - if executor is None or executor._watchdog_stop or executor._ready is None: + if executor is None or executor._watchdog_stop or not executor._flag_sync: return try: executor._watchdog_tick(suspects) diff --git a/tests/engine/test_cpu_moe_graph_safety.py b/tests/engine/test_cpu_moe_graph_safety.py new file mode 100644 index 000000000..d05de7efc --- /dev/null +++ b/tests/engine/test_cpu_moe_graph_safety.py @@ -0,0 +1,155 @@ +from types import SimpleNamespace + +import pytest + + +@pytest.mark.parametrize("backend", ["cpu", "hybrid", "offload"]) +def test_rocm_unsafe_cpu_executor_disables_graph(monkeypatch, backend): + """The same gate covers cpu/hybrid and offload + moe_cpu_layers.""" + import freetoken.engine.engine as engine + + monkeypatch.setattr(engine, "is_rocm", lambda: True) + config = SimpleNamespace( + moe_backend=backend, + moe_cpu_layers="0" if backend == "offload" else None, + cuda_graph_bs=None, + cuda_graph_max_bs=32, + ) + executor = SimpleNamespace(graph_capture_safe=False) + + assert engine._disable_unsafe_rocm_cpu_moe_graph(config, executor) + assert config.cuda_graph_bs == [] + assert config.cuda_graph_max_bs == 0 + + +def test_rocm_verified_cpu_executor_keeps_graph(monkeypatch): + import freetoken.engine.engine as engine + + monkeypatch.setattr(engine, "is_rocm", lambda: True) + config = SimpleNamespace(cuda_graph_bs=[1, 2, 4], cuda_graph_max_bs=4) + + assert not engine._disable_unsafe_rocm_cpu_moe_graph( + config, SimpleNamespace(graph_capture_safe=True) + ) + assert config.cuda_graph_bs == [1, 2, 4] + assert config.cuda_graph_max_bs == 4 + + +def test_non_rocm_keeps_existing_cuda_host_callback_fallback(monkeypatch): + import freetoken.engine.engine as engine + + monkeypatch.setattr(engine, "is_rocm", lambda: False) + config = SimpleNamespace(cuda_graph_bs=None, cuda_graph_max_bs=16) + + assert not engine._disable_unsafe_rocm_cpu_moe_graph( + config, SimpleNamespace(graph_capture_safe=False) + ) + assert config.cuda_graph_bs is None + assert config.cuda_graph_max_bs == 16 + + +@pytest.mark.parametrize( + ("cuda_graph_bs", "cuda_graph_max_bs", "expected"), + [ + (None, 32, False), + ([1, 2], 2, False), + ([1], 0, False), + ([], 32, True), + (None, 0, True), + ], +) +def test_rocm_auto_hybrid_requires_graph_to_be_disabled( + monkeypatch, cuda_graph_bs, cuda_graph_max_bs, expected +): + import freetoken.engine.engine as engine + + monkeypatch.setattr(engine, "is_rocm", lambda: True) + config = SimpleNamespace( + cuda_graph_bs=cuda_graph_bs, cuda_graph_max_bs=cuda_graph_max_bs + ) + + assert engine._auto_hybrid_allowed_before_rocm_probe(config) is expected + + +def test_cpu_moe_flag_slots_cover_every_graph_batch_size(): + import freetoken.engine.engine as engine + + config = SimpleNamespace(cuda_graph_bs=None, cuda_graph_max_bs=128) + + # [1, 2, 4] plus 8..128 in steps of 8: 19 distinct captured sizes. + assert engine._cpu_moe_flag_slots_per_layer(config, free_memory=0) == 19 + + +def test_cpu_moe_flag_slots_keep_eager_headroom_for_small_graph_sets(): + import freetoken.engine.engine as engine + + config = SimpleNamespace(cuda_graph_bs=[1, 2, 2, 4], cuda_graph_max_bs=128) + + assert engine._cpu_moe_flag_slots_per_layer(config, free_memory=0) == 16 + + +def test_cuda_flag_memops_keep_module_level_api(monkeypatch): + """The ownership refactor must not reroute CUDA through the ROCm-only methods.""" + import freetoken.moe.cpu_executor as cpu_executor + + calls = [] + + class TensorStub: + shape = (1, 8) + + def copy_(self, *args, **kwargs): + return self + + def to(self, *args, **kwargs): + return self + + executor = object.__new__(cpu_executor.CpuMoeExecutor) + executor._flag_sync = True + executor._shared_flag_signal = False + executor._gpu_prequant = False + executor._flag_slots = {(0, 1): 3} + executor._cpu_moe = SimpleNamespace( + memop_submit=lambda *args: calls.append(("submit", args)), + memop_sync=lambda *args: calls.append(("sync", args)), + ) + executor._ext = SimpleNamespace( + submit_flag_memop=lambda *args: pytest.fail("ROCm submit used on CUDA"), + sync_flag_memop=lambda *args: pytest.fail("ROCm sync used on CUDA"), + ) + executor._ready = SimpleNamespace(data_ptr=lambda: 101) + executor._done = SimpleNamespace(data_ptr=lambda: 202) + io = {name: TensorStub() for name in ("x", "ids", "w", "y")} + executor._io_for = lambda bs: io + executor._task_for = lambda layer_id, bs: 303 + executor._io = {1: io} + + monkeypatch.setattr(cpu_executor, "_IS_ROCM", False) + monkeypatch.setattr( + cpu_executor.torch.cuda, + "current_stream", + lambda: SimpleNamespace(cuda_stream=404), + ) + monkeypatch.setattr(cpu_executor.torch, "empty_like", lambda value: TensorStub()) + + tensor = TensorStub() + pending = executor.decode_submit(0, tensor, tensor, tensor) + executor.decode_sync(pending) + + assert calls == [ + ("submit", (404, 202, 101, 3)), + ("sync", (404, 202, 3)), + ] + + +def test_direct_rocm_capture_without_flag_sync_fails_closed(monkeypatch): + import freetoken.moe.cpu_executor as cpu_executor + + monkeypatch.setattr(cpu_executor, "_IS_ROCM", True) + monkeypatch.setattr( + cpu_executor.torch.cuda, "is_current_stream_capturing", lambda: True + ) + executor = object.__new__(cpu_executor.CpuMoeExecutor) + executor._flag_sync = False + + with pytest.raises(RuntimeError, match="not CUDA-graph safe on ROCm"): + executor.decode_submit(0, None, None, None) diff --git a/tests/moe/test_cpu_moe.py b/tests/moe/test_cpu_moe.py index f496de314..5b58c7d26 100644 --- a/tests/moe/test_cpu_moe.py +++ b/tests/moe/test_cpu_moe.py @@ -4,9 +4,9 @@ GPU decode kernels (bf16/nvfp4/mxfp4/ds_fp4) on identical banks and routing across batch sizes (fp32-accumulate, so the only spread is reduction order -> tight tol). -Part 2 -- CUDA-graph capture/replay: the cudaLaunchHostFunc submit/sync host nodes -end to end (eager decode, capture, replay with *new* data -> result tracks the new -routing, i.e. the dependency flows through the pinned buffers, not baked-in). +Part 2 -- CUDA-graph capture/replay: the native flag handshake end to end (eager +decode, capture, replay with *new* data -> result tracks the new routing, i.e. the +dependency flows through the pinned buffers, not baked-in). """ from __future__ import annotations @@ -502,15 +502,136 @@ def test_cpu_moe_decode_cuda_graph_replay(): assert (layer, bs) in ex._flag_slots, "flag slot expected for the decode task" slot = ex._flag_slots[(layer, bs)] assert ex._ext.flag_served_count(slot) >= 4, "1 eager + 3 replay dispatches expected" - assert int(ex._done[slot]) == 1 and int(ex._ready[slot]) == 0, "handshake at rest" + assert ex._ext.flag_done_value(slot) == 1, "completion flag at rest" + assert ex._ext.flag_ready_value(slot) == 0, "doorbell must be consumed" assert int(ex._err.sum()) == 0, "watchdog must not fire in normal operation" ex.raise_if_unhealthy() print("cpu moe cuda graph replay OK") +@pytest.mark.skipif(getattr(torch.version, "hip", None) is None, reason="ROCm guard") +def test_rocm_cpu_moe_capture_without_flag_sync_fails_closed(monkeypatch): + """ROCm must never capture the known-unsafe hipLaunchHostFunc fallback.""" + import freetoken.moe.cpu_executor as cpu_executor + + monkeypatch.setattr(cpu_executor, "_FLAG_SYNC", False) + cache = _make_cache(1, 4, 64, 32) + dev = torch.device("cuda") + stream = torch.cuda.Stream() + torch.cuda.set_stream(stream) + executor = cpu_executor.CpuMoeExecutor( + cache, + top_k=2, + activation="silu", + apply_router_weight_on_input=False, + num_threads=2, + max_tokens=1, + device=dev, + ) + hidden = torch.randn(1, 64, device=dev, dtype=torch.bfloat16) + ids = torch.randint(0, 4, (1, 2), device=dev, dtype=torch.int32) + weights = torch.rand(1, 2, device=dev, dtype=torch.float32) + + with pytest.raises(RuntimeError, match="not CUDA-graph safe on ROCm"): + with torch.cuda.graph(torch.cuda.CUDAGraph(), stream=stream): + hidden.add_(0) # keep capture valid when the executor rejects its own nodes + executor.decode(0, hidden, weights, ids) + + +@pytest.mark.skipif(getattr(torch.version, "hip", None) is None, reason="ROCm guard") +def test_rocm_graph_memop_param_storage_is_bounded_across_rebuilds(): + """Rebuilding graphs must reuse executor-owned HIP batch-op parameters. + + ROCm 7.14 retains each graph node's ``paramArray`` pointer. The executor keeps + one immutable submit array per flag slot plus one shared wait array, so graph + rebuilds must neither invalidate those addresses nor append new storage. + """ + from freetoken.moe.cpu_executor import CpuMoeExecutor + + torch.manual_seed(37) + L, E, H, I, top_k, bs = 1, 4, 64, 32, 2, 1 + cache = _make_cache(L, E, H, I) + stream = torch.cuda.Stream() + torch.cuda.set_stream(stream) + executor = CpuMoeExecutor( + cache, + top_k=top_k, + activation="silu", + apply_router_weight_on_input=False, + num_threads=2, + max_tokens=bs, + device=torch.device("cuda"), + ) + if not executor._flag_sync: + pytest.skip("native ROCm graph handshake unavailable") + + hidden = torch.randn(bs, H, device="cuda", dtype=torch.bfloat16) + ids = torch.randint(0, E, (bs, top_k), device="cuda", dtype=torch.int32) + weights = torch.rand(bs, top_k, device="cuda", dtype=torch.float32) + executor.decode(0, hidden, weights, ids) # materialize the task and flag slot + torch.cuda.synchronize() + + expected_params = 2 * executor._flag_capacity + 1 + assert executor._ext.graph_memop_param_count() == expected_params + for _ in range(16): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = executor.decode(0, hidden, weights, ids) + graph.replay() + torch.cuda.synchronize() + assert torch.isfinite(out).all() + assert executor._ext.graph_memop_param_count() == expected_params + del graph, out + + +@pytest.mark.skipif(getattr(torch.version, "hip", None) is None, reason="ROCm guard") +def test_rocm_graph_memops_support_more_than_sixteen_batch_sizes(): + """Every configured graph size gets a native slot, including size sets > 16.""" + from freetoken.moe.cpu_executor import CpuMoeExecutor + + torch.manual_seed(41) + batch_sizes = list(range(1, 20)) + L, E, H, I, top_k = 1, 4, 64, 32, 2 + cache = _make_cache(L, E, H, I) + stream = torch.cuda.Stream() + torch.cuda.set_stream(stream) + executor = CpuMoeExecutor( + cache, + top_k=top_k, + activation="silu", + apply_router_weight_on_input=False, + num_threads=2, + max_tokens=max(batch_sizes), + device=torch.device("cuda"), + flag_slots_per_layer=len(batch_sizes), + ) + if not executor._flag_sync: + pytest.skip("native ROCm graph handshake unavailable") + + for bs in batch_sizes: + hidden = torch.randn(bs, H, device="cuda", dtype=torch.bfloat16) + ids = torch.randint(0, E, (bs, top_k), device="cuda", dtype=torch.int32) + weights = torch.rand(bs, top_k, device="cuda", dtype=torch.float32) + executor.decode(0, hidden, weights, ids) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph, stream=stream): + out = executor.decode(0, hidden, weights, ids) + graph.replay() + torch.cuda.synchronize() + slot = executor._flag_slots[(0, bs)] + assert executor._ext.flag_served_count(slot) >= 2 + assert torch.isfinite(out).all() + del graph, out + + assert len(executor._flag_slots) == len(batch_sizes) + assert executor._ext.graph_memop_param_count() == 2 * len(batch_sizes) + 1 + + def test_cpu_moe_decode_cuda_graph_replay_mxfp4(): - """gpt-oss mxfp4 path under capture/replay: the host nodes must recompute the + """gpt-oss mxfp4 path under capture/replay: the handshake must recompute the clamped-swiglu+bias GEMV from the freshly written pinned routing on each replay.""" from freetoken.moe.cpu_executor import CpuMoeExecutor from freetoken.moe.fused_mxfp4 import run_mxfp4_splitk_decode_experts as _run_mxfp4_splitk_decode_experts diff --git a/tests/moe/test_cpu_moe_q4_0.py b/tests/moe/test_cpu_moe_q4_0.py index b5ec2f09b..1e7c63cd7 100644 --- a/tests/moe/test_cpu_moe_q4_0.py +++ b/tests/moe/test_cpu_moe_q4_0.py @@ -7,8 +7,8 @@ production bf16 GPU decode kernel on byte-identical banks: both are W4A16, so the only spread is weight bf16-rounding + reduction order -> tight relative tolerance. -Part 2 covers CUDA-graph capture/replay (the cudaLaunchHostFunc submit/sync nodes -must recompute from the freshly written pinned routing on each replay). +Part 2 covers CUDA-graph capture/replay (the native flag handshake must recompute +from the freshly written pinned routing on each replay). """ from __future__ import annotations @@ -148,7 +148,7 @@ def test_cpu_decode_q4_0_matches_ggml_mmvq(): def test_cpu_moe_decode_q4_0_cuda_graph_replay(): - """Q4_0 CPU path under capture/replay: the host nodes must recompute the GEMV from + """Q4_0 CPU path under capture/replay: the handshake must recompute the GEMV from the freshly written pinned routing on each replay (dep flows through pinned buffers).""" from freetoken.moe.cpu_executor import CpuMoeExecutor from freetoken.moe.fused import fused_experts_decode_impl